remove max new tokens

KonferCA · Nov 24, 2024 · 1d87be6 · 1d87be6
1 parent 5af4c90
commit 1d87be6
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 3 deletions.
diff --git a/backend/generate_answer.py b/backend/generate_answer.py
@@ -14,8 +14,7 @@ def load_base_model():
 
     # Load model in 8-bit to reduce memory usage
     base_model = AutoModelForCausalLM.from_pretrained(Config.MODEL_NAME, 
-    token=Config.HUGGINGFACE_ACCESS_TOKEN,
-    max_new_tokens=8096)
+    token=Config.HUGGINGFACE_ACCESS_TOKEN)
 
     return base_model, tokenizer
 

diff --git a/backend/lora.py b/backend/lora.py
@@ -117,7 +117,6 @@ def apply_lora_weights_to_model(base_model_name, lora_weights_dir):
         device_map="auto",
         quantization_config=bnb_config,
         torch_dtype=torch.float16,
-        max_new_tokens=8096
     )
 
     model = PeftModel.from_pretrained(base_model, lora_weights_dir)