Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from ctransformers import AutoModelForCausalLM
+from llama_cpp import Llama
 import os
 
 # Download and load the GGUF model
@@ -14,27 +14,15 @@ if not os.path.exists(model_path):
     print("Model downloaded!")
 
 # Load the model
-llm = AutoModelForCausalLM.from_pretrained(
-    model_path,
-    model_type="llama",
-    context_length=2048,
-    gpu_layers=0  # Set higher if GPU available
+print("Loading model...")
+llm = Llama(
+    model_path=model_path,
+    n_ctx=2048,
+    n_threads=4,
+    n_gpu_layers=0,
+    verbose=False
 )
-
-def format_prompt(message, history):
-    """Format the conversation into Llama 3.2 chat format"""
-    prompt = ""
-    
-    # Add chat history
-    for user_msg, bot_msg in history:
-        prompt += f"<|start_header_id|>user<|end_header_id|>\n\n{user_msg}<|eot_id|>"
-        prompt += f"<|start_header_id|>assistant<|end_header_id|>\n\n{bot_msg}<|eot_id|>"
-    
-    # Add current message
-    prompt += f"<|start_header_id|>user<|end_header_id|>\n\n{message}<|eot_id|>"
-    prompt += "<|start_header_id|>assistant<|end_header_id|>\n\n"
-    
-    return prompt
+print("Model loaded!")
 
 def chat(message, history):
     """
@@ -44,19 +32,27 @@ def chat(message, history):
         message: Current user message
         history: List of [user_msg, bot_msg] pairs
     """
-    # Format the prompt
-    prompt = format_prompt(message, history)
+    # Build conversation with proper Llama format
+    messages = []
+    
+    # Add chat history
+    for user_msg, bot_msg in history:
+        messages.append({"role": "user", "content": user_msg})
+        messages.append({"role": "assistant", "content": bot_msg})
+    
+    # Add current message
+    messages.append({"role": "user", "content": message})
     
     # Generate response
-    response = llm(
-        prompt,
-        max_new_tokens=512,
+    response = llm.create_chat_completion(
+        messages=messages,
+        max_tokens=512,
         temperature=0.7,
         top_p=0.9,
-        stop=["<|eot_id|>", "<|start_header_id|>"]
     )
     
-    return response.strip()
+    # Extract the assistant's response
+    return response["choices"][0]["message"]["content"]
 
 # Create Gradio interface
 demo = gr.ChatInterface(