Fixed Llama3.1 support

eisenzopf · Jan 20, 2025 · 07ca7fe · 07ca7fe
1 parent f0248fe
commit 07ca7fe
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/lapet/llama.py b/lapet/llama.py
@@ -39,13 +39,14 @@ def post_process_output(self, prompt, output):
 
 class Llama31ModelHandler():
     def load_model_and_tokenizer(self, device, model_id):
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
 
         model = AutoModelForCausalLM.from_pretrained(
             model_id,
-            device_map=device
+            device_map=device,
+            trust_remote_code=True
         )
 
         print(model_id + " loaded.")

diff --git a/pyproject.toml b/pyproject.toml
@@ -6,7 +6,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "lapet"
-version = "0.8.4"
+version = "0.8.5"
 authors = [
     {name = "Jonathan Eisenzopf", email = "[email protected]"},
 ]