Update app.py

app.py

@@ -3,6 +3,7 @@ from transformers import AutoTokenizer, LlamaForCausalLM, BitsAndBytesConfig
 from transformers.generation.stopping_criteria import StoppingCriteria, StoppingCriteriaList
 from peft import PeftModel
 import gradio as gr
+import os
 
 # Add this new class for custom stopping criteria
 class SentenceEndingCriteria(StoppingCriteria):
@@ -15,24 +16,25 @@ class SentenceEndingCriteria(StoppingCriteria):
         return last_token in self.end_tokens
 
 def load_model():
-    # Modify the model path to use the Hugging Face model ID
-    model_path = "Cioni223/mymodel"  # Replace with your actual model path on HF
+    model_path = "Cioni223/mymodel"
+    token = os.environ.get("HUGGINGFACE_TOKEN")  # Ensure you set this environment variable
+
     tokenizer = AutoTokenizer.from_pretrained(
         model_path,
         use_fast=False,
         padding_side="left",
         model_max_length=4096,
-        token=True  # Add this if your model is private
+        token=token
     )
     
     tokenizer.pad_token = tokenizer.eos_token
     
-    # Load merged model with quantization
     model = LlamaForCausalLM.from_pretrained(
         model_path,
         device_map="auto",
         torch_dtype=torch.float16,
-        quantization_config=BitsAndBytesConfig(load_in_8bit=True)
+        quantization_config=BitsAndBytesConfig(load_in_8bit=True),
+        use_auth_token=token
     )
     
     return model, tokenizer