Update app.py

app.py

@@ -1,64 +1,128 @@
+import torch
+from transformers import AutoTokenizer, LlamaForCausalLM, BitsAndBytesConfig
+from transformers.generation.stopping_criteria import StoppingCriteria, StoppingCriteriaList
+from peft import PeftModel
 import gradio as gr
-from huggingface_hub import InferenceClient
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("Cioni223/mymodel")
+# 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
+    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
+    )
+    
+    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)
+    )
+    
+    return model, tokenizer
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
+def format_chat_history(history):
+    formatted_history = ""
+    for user_msg, assistant_msg in history:
+        if user_msg:
+            formatted_history += f"<|start_header_id|>user<|end_header_id|>{user_msg}<|eot_id|>\n"
+        if assistant_msg:
+            formatted_history += f"<|start_header_id|>assistant<|end_header_id|>{assistant_msg}<|eot_id|>\n"
+    return formatted_history
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+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|>"""
+    
+    chat_history = format_chat_history(history)
+    
+    formatted_prompt = f"""{system_prompt}
+{chat_history}<|start_header_id|>user<|end_header_id|>{message}<|eot_id|>
+<|start_header_id|>assistant<|end_header_id|>"""
+    
+    inputs = tokenizer(
+        formatted_prompt,
+        return_tensors="pt",
+        padding=True
+    ).to(model.device)
+    
+    # Create stopping criteria
+    end_tokens = [
+        tokenizer.encode(".")[0],
+        tokenizer.encode("!")[0],
+        tokenizer.encode("?")[0],
+        tokenizer.encode("<|eot_id|>", add_special_tokens=False)[0]
+    ]
+    stopping_criteria = StoppingCriteriaList([
+        SentenceEndingCriteria(tokenizer, end_tokens)
+    ])
+    
+    # Modified generation parameters
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=300,
+            temperature=0.4,
+            do_sample=True,
+            top_p=0.95,
+            top_k=50,
+            repetition_penalty=1.2,
+            no_repeat_ngram_size=3,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.encode("<|eot_id|>", add_special_tokens=False)[0],
+            stopping_criteria=stopping_criteria
+        )
+    
+    response = tokenizer.decode(outputs[0], skip_special_tokens=False)
+    
+    try:
+        assistant_parts = response.split("<|start_header_id|>assistant<|end_header_id|>")
+        last_response = assistant_parts[-1].split("<|eot_id|>")[0].strip()
+        
+        # Ensure response ends with proper punctuation
+        if not any(last_response.rstrip().endswith(punct) for punct in ['.', '!', '?']):
+            # Find the last complete sentence
+            sentences = last_response.split('.')
+            if len(sentences) > 1:
+                last_response = '.'.join(sentences[:-1]) + '.'
+            
+        return last_response
+    except:
+        return "I apologize, but I couldn't generate a proper response. Please try again."
 
-    messages.append({"role": "user", "content": message})
+# Load model and tokenizer
+print("Loading model...")
+model, tokenizer = load_model()
+print("Model loaded!")
 
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-
-        response += token
-        yield response
-
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
+# Create Gradio interface with chat
 demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
+    fn=chat_response,
+    title="Admissions Agent Assistant",
+    description="Chat with an AI-powered admissions coordinator. The agent will maintain context of your conversation.",
+    examples=[
+        "I need help with addiction treatment",
+        "What insurance do you accept?",
+        "How long are your treatment programs?",
+        "Can you help with mental health issues?"
+    ]
 )
 
-
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch() # Remove share=True as it's not needed for HF Spaces