Update app.py

app.py

@@ -7,6 +7,12 @@ 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():
@@ -23,7 +29,6 @@ def load_model():
     
     tokenizer.pad_token = tokenizer.eos_token
     
-
     model = LlamaForCausalLM.from_pretrained(
         model_path,
         device_map="auto",
@@ -43,7 +48,6 @@ def format_chat_history(history):
             formatted_history += f"<|start_header_id|>assistant<|end_header_id|>{assistant_msg}<|eot_id|>\n"
     return formatted_history
 
-
 def chat_response(message, history):
     # Format the prompt with system message and chat history
     system_prompt = """<|start_header_id|>system<|end_header_id|>You are Fred, a virtual admissions coordinator for Haven Health Management, a mental health and substance abuse treatment facility. Your role is to respond conversationally and empathetically, like a human agent, using 1-2 sentences per response while guiding the conversation effectively. Your primary goal is to understand the caller's reason for reaching out, gather their medical history, and obtain their insurance details, ensuring the conversation feels natural and supportive. Once all the information is gathered politely end the conversation and if the user is qualified tell the user a live agent will reach out soon. Note: Medicaid is not accepted as insurance.<|eot_id|>"""
@@ -122,6 +126,5 @@ demo = gr.ChatInterface(
     ]
 )
 
-
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch() # Remove share=True as it's not needed for HF Spaces