Skip to content

Commit e51acb0

Browse files
authored
Add "trust_remote_code" to Sentence Transformer snippets (#936)
Hello! ## Pull Request overview * Add "trust_remote_code" to Sentence Transformer snippets ## Details The `SentenceTransformer` class supports the `trust_remote_code` argument, so let's take advantage of it. I didn't test it with `moon-landing`, etc. Some examples where this would be used: * https://huggingface.co/nomic-ai/nomic-embed-text-v1.5 * https://huggingface.co/jinaai/jina-embeddings-v3 --- - Tom Aarsen
1 parent c9d9930 commit e51acb0

File tree

1 file changed

+8
-4
lines changed

1 file changed

+8
-4
lines changed

packages/tasks/src/model-libraries-snippets.ts

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -678,11 +678,15 @@ export const sampleFactory = (model: ModelData): string[] => [
678678
`python -m sample_factory.huggingface.load_from_hub -r ${model.id} -d ./train_dir`,
679679
];
680680

681-
export const sentenceTransformers = (model: ModelData): string[] => [
682-
`from sentence_transformers import SentenceTransformer
681+
export const sentenceTransformers = (model: ModelData): string[] => {
682+
const remote_code_snippet = model.tags.includes(TAG_CUSTOM_CODE) ? ", trust_remote_code=True" : "";
683683

684-
model = SentenceTransformer("${model.id}")`,
685-
];
684+
return [
685+
`from sentence_transformers import SentenceTransformer
686+
687+
model = SentenceTransformer("${model.id}"${remote_code_snippet})`,
688+
];
689+
};
686690

687691
export const setfit = (model: ModelData): string[] => [
688692
`from setfit import SetFitModel

0 commit comments

Comments
 (0)