Fixing Llama3.1 support

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

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/lapet/llama.py b/lapet/llama.py
@@ -1,6 +1,6 @@
 import re
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
 
 from .handler import ModelHandler
 class Llama3ModelHandler():
@@ -43,8 +43,18 @@ def load_model_and_tokenizer(self, device, model_id):
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
 
+        # Load the model's config first
+        config = AutoConfig.from_pretrained(model_id)
+
+        # Override only the required RoPE parameters
+        #config.rope_scaling = {
+        #    "type": "dynamic",
+        #    "factor": 8.0
+        #}
+
         model = AutoModelForCausalLM.from_pretrained(
             model_id,
+            config=config,
             device_map=device,
             trust_remote_code=True
         )

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