Update app.py

app.py

@@ -3,37 +3,33 @@ 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):
-    def __init__(self, tokenizer, end_tokens):
-        self.tokenizer = tokenizer
-        self.end_tokens = end_tokens
-        
-    def __call__(self, input_ids, scores, **kwargs):
-        last_token = input_ids[0][-1]
         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