Make attn_implementation configurable for huggingface models (#146)

sillsdev · Jan 14, 2025 · 6de035f · 6de035f
1 parent 9110963
commit 6de035f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/machine/jobs/huggingface/hugging_face_nmt_model_factory.py b/machine/jobs/huggingface/hugging_face_nmt_model_factory.py
@@ -65,7 +65,11 @@ def init(self) -> None:
         )
         self._model = cast(
             PreTrainedModel,
-            AutoModelForSeq2SeqLM.from_pretrained(self._config.huggingface.parent_model_name, config=config),
+            AutoModelForSeq2SeqLM.from_pretrained(
+                self._config.huggingface.parent_model_name,
+                config=config,
+                attn_implementation=self._config.huggingface.attn_implementation,
+            ),
         )
 
     def create_source_tokenizer_trainer(self, corpus: TextCorpus) -> Trainer:

diff --git a/machine/jobs/settings.yaml b/machine/jobs/settings.yaml
@@ -27,6 +27,7 @@ default:
     tokenizer:
       add_unk_src_tokens: true
       add_unk_trg_tokens: true
+    attn_implementation: sdpa
   thot_mt:
     word_alignment_model_type: hmm
     tokenizer: latin