Update app.py

app.py

@@ -1,43 +1,78 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-# MODEL_ID can be any HF model repo you want to load
-MODEL_ID = "unsloth/Llama-3.2-1B-Instruct"  # replace with your own on the Hub
+# Load model and tokenizer
+model_name = "unsloth/Llama-3.2-1B-Instruct"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    device_map="auto"
+)
 
-# load tokenizer + model
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
-
-# chat function
-def chat(input_text, history=[]):
-    # append user input to history
-    history = history or []
-    history.append(("You", input_text))
-
-    # encode input
-    inputs = tokenizer(
-        input_text + tokenizer.eos_token,
-        return_tensors="pt"
+def chat(message, history):
+    """
+    Process chat messages and generate responses.
+    
+    Args:
+        message: Current user message
+        history: List of [user_msg, bot_msg] pairs
+    """
+    # 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})
+    
+    # Apply chat template
+    prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
     )
-
-    # generate a response
-    output_ids = model.generate(
+    
+    # Tokenize
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    
+    # Generate response
+    outputs = model.generate(
         **inputs,
-        max_new_tokens=100,
+        max_new_tokens=512,
+        temperature=0.7,
+        top_p=0.9,
+        do_sample=True,
         pad_token_id=tokenizer.eos_token_id
     )
+    
+    # Decode and extract only the new response
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # Extract only the assistant's response (after the last prompt)
+    response = response.split("assistant")[-1].strip()
+    
+    return response
 
-    response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    history.append(("Bot", response))
-
-    # return formatted history
-    return "", history
-
-# Gradio UI
-with gr.Blocks() as demo:
-    chatbot = gr.Chatbot()
-    user_input = gr.Textbox(placeholder="Type a message…")
-    user_input.submit(chat, [user_input, chatbot], [user_input, chatbot])
+# Create Gradio interface
+demo = gr.ChatInterface(
+    fn=chat,
+    title="Llama 3.2 1B Instruct Chatbot",
+    description="Chat with Llama 3.2 1B Instruct model. Ask me anything!",
+    examples=[
+        "What is artificial intelligence?",
+        "Write a short poem about coding",
+        "Explain quantum computing in simple terms"
+    ],
+    theme=gr.themes.Soft(),
+    retry_btn="Retry",
+    undo_btn="Undo",
+    clear_btn="Clear"
+)
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()