Update app.py

app.py

@@ -38,20 +38,21 @@ except OSError as e:
 
 
 # --- Chat Function ---
-
 def chat_with_llm(prompt, history):
-    """Generates a response from the LLM."""
-
+    """Generates a response from the LLM, handling history correctly."""
     formatted_prompt = ""
     if history:
-        for user_msg, ai_msg in history:
-            formatted_prompt += f"{tokenizer.bos_token}{user_msg}{tokenizer.eos_token}"
-            formatted_prompt += f"{ai_msg}{tokenizer.eos_token}"
+        for item in history:
+            if item["role"] == "user":
+                formatted_prompt += f"{tokenizer.bos_token}{item['content']}{tokenizer.eos_token}"
+            elif item["role"] == "assistant":
+                formatted_prompt += f"{item['content']}{tokenizer.eos_token}"
 
     formatted_prompt += f"{tokenizer.bos_token}{prompt}{tokenizer.eos_token}"
+
     try:
-      pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device_map="auto")
-      result = pipe(
+        pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device_map="auto")
+        result = pipe(
             formatted_prompt,
             max_new_tokens=256,
             do_sample=True,
@@ -61,19 +62,20 @@ def chat_with_llm(prompt, history):
             return_full_text=False,
             pad_token_id=tokenizer.eos_token_id,
         )
-      response = result[0]['generated_text'].strip()
-      return response
+        response = result[0]['generated_text'].strip()
+        return response
     except Exception as e:
         return f"Error during generation: {e}"
 
+
+
 # --- Gradio Interface ---
 # Use the 'messages' format for chatbot
 def predict(message, history):
     history = history or []
     response = chat_with_llm(message, history)
-    # history.append((message, response)) # Old format
-    history.append({"role": "user", "content": message}) # Append user message in 'messages' format
-    history.append({"role": "assistant", "content": response})  # Append assistant (AI) message
+    history.append({"role": "user", "content": message})
+    history.append({"role": "assistant", "content": response})
     return "", history
 
 with gr.Blocks() as demo:
@@ -83,9 +85,7 @@ with gr.Blocks() as demo:
     msg = gr.Textbox(label="Your Message", placeholder="Type your message here...")
     clear = gr.Button("Clear")
 
-
     msg.submit(predict, [msg, chatbot], [msg, chatbot])
-     # Clear needs to be updated for the messages format
-    clear.click(lambda: [], [], chatbot, queue=False) # Return empty list for history
+    clear.click(lambda: [], [], chatbot, queue=False)
 
 demo.launch(share=True)