Update app.py

app.py

@@ -4,12 +4,16 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStream
 from peft import PeftModel
 from threading import Thread
 
-# Your model configuration
+# ------------------------------------------------------------------
+# 1. Model setup
+# ------------------------------------------------------------------
 BASE_MODEL = "unsloth/Qwen3-4B-Instruct-2507"
 LORA_MODEL = "michsethowusu/twi_code_assistant"
 
-print("Loading base model...")
+print("Loading tokenizer…")
 tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+
+print("Loading base model…")
 base_model = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL,
     torch_dtype=torch.float16,
@@ -17,131 +21,115 @@ base_model = AutoModelForCausalLM.from_pretrained(
     low_cpu_mem_usage=True
 )
 
-print("Loading LoRA adapters...")
+print("Loading LoRA adapters…")
 model = PeftModel.from_pretrained(base_model, LORA_MODEL)
-model = model.merge_and_unload()  # Merge for faster inference
+model = model.merge_and_unload()
 print("Model ready!")
 
-def generate_response(message, history, temperature, top_p, top_k, max_tokens):
-    """Generate response from the model with streaming"""
-    
-    # Build conversation history - history is list of [user_msg, bot_msg] pairs
-    messages = []
-    for user_msg, assistant_msg in history:
-        if user_msg:
-            messages.append({"role": "user", "content": user_msg})
-        if assistant_msg:
-            messages.append({"role": "assistant", "content": assistant_msg})
-    
-    # Add current message
-    messages.append({"role": "user", "content": message})
-    
-    # Apply chat template
-    text = tokenizer.apply_chat_template(
+# ------------------------------------------------------------------
+# 2. Generation helper
+# ------------------------------------------------------------------
+def generate_response(message: str, history: list[dict], temperature, top_p, top_k, max_tokens):
+    """
+    message: str – the newest user message
+    history: list[dict] – previous turns in {"role": "user"|"assistant", "content": "…"} format
+    yields partial assistant reply strings
+    """
+    messages = history + [{"role": "user", "content": message}]
+
+    prompt = tokenizer.apply_chat_template(
         messages,
         tokenize=False,
         add_generation_prompt=True
     )
-    
-    # Tokenize
-    inputs = tokenizer(text, return_tensors="pt").to(model.device)
-    
-    # Setup streaming
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
     streamer = TextIteratorStreamer(
         tokenizer,
         skip_prompt=True,
         skip_special_tokens=True
     )
-    
-    # Generation kwargs
-    generation_kwargs = {
+
+    gen_kwargs = dict(
         **inputs,
-        "max_new_tokens": max_tokens,
-        "temperature": temperature,
-        "top_p": top_p,
-        "top_k": top_k,
-        "do_sample": True,
-        "streamer": streamer,
-    }
-    
-    # Start generation in separate thread
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+        max_new_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        do_sample=True,
+        streamer=streamer,
+    )
+
+    thread = Thread(target=model.generate, kwargs=gen_kwargs)
     thread.start()
-    
-    # Stream the response
-    partial_message = ""
+
+    partial = ""
     for new_text in streamer:
-        partial_message += new_text
-        yield partial_message
-    
+        partial += new_text
+        yield partial
+
     thread.join()
 
-# Create Gradio interface
+# ------------------------------------------------------------------
+# 3. Gradio event helpers
+# ------------------------------------------------------------------
+def user_submit(user_message, history):
+    # history is list[dict] – append user message
+    return "", history + [{"role": "user", "content": user_message}]
+
+
+def bot_respond(history, temperature, top_p, top_k, max_tokens):
+    user_turn = history[-1]["content"]
+    history_before = history[:-1]
+
+    assistant_text = ""
+    for assistant_text in generate_response(
+        user_turn, history_before, temperature, top_p, top_k, max_tokens
+    ):
+        # update last dict incrementally
+        history[-1] = {"role": "assistant", "content": assistant_text}
+        yield history
+
+# ------------------------------------------------------------------
+# 4. Gradio UI
+# ------------------------------------------------------------------
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
         # 🇬🇭 Twi Code Assistant
-        
-        A fine-tuned Qwen3-4B model specialized for coding assistance in Twi language context.
+        A fine-tuned Qwen3-4B model specialised for coding assistance in Twi language context.
         Ask me anything about programming, and I'll help you out!
         """
     )
-    
+
     chatbot = gr.Chatbot(
         height=500,
         label="Chat History",
-        avatar_images=(None, "https://em-content.zobj.net/source/twitter/53/robot-face_1f916.png")
+        type="messages",  # NEW FORMAT
+        avatar_images=(None, "https://em-content.zobj.net/source/twitter/53/robot-face_1f916.png"),
     )
-    
+
     with gr.Row():
         msg = gr.Textbox(
             label="Your Message",
-            placeholder="Ask me a coding question...",
+            placeholder="Ask me a coding question…",
             scale=4,
-            lines=2
+            lines=2,
         )
         submit = gr.Button("Send 🚀", scale=1, variant="primary")
-    
+
     with gr.Accordion("⚙️ Generation Parameters", open=False):
         gr.Markdown("*Adjust these settings to control the response style*")
-        
-        temperature = gr.Slider(
-            minimum=0.1,
-            maximum=2.0,
-            value=0.7,
-            step=0.1,
-            label="Temperature",
-            info="Higher = more creative, Lower = more focused"
-        )
-        top_p = gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.8,
-            step=0.05,
-            label="Top P",
-            info="Nucleus sampling threshold"
-        )
-        top_k = gr.Slider(
-            minimum=1,
-            maximum=100,
-            value=20,
-            step=1,
-            label="Top K",
-            info="Number of top tokens to consider"
-        )
-        max_tokens = gr.Slider(
-            minimum=64,
-            maximum=2048,
-            value=512,
-            step=64,
-            label="Max Tokens",
-            info="Maximum length of response"
-        )
-    
-    with gr.Row():
-        clear = gr.Button("🗑️ Clear Chat")
-        
-    # Example prompts
+        temperature = gr.Slider(0.1, 2.0, 0.7, step=0.1, label="Temperature")
+        top_p = gr.Slider(0.1, 1.0, 0.8, step=0.05, label="Top P")
+        top_k = gr.Slider(1, 100, 20, step=1, label="Top K")
+        max_tokens = gr.Slider(64, 2048, 512, step=64, label="Max Tokens")
+
+    clear = gr.Button("🗑️ Clear Chat")
+
+    # ------------------------------------------------------------------
+    # 5. Examples
+    # ------------------------------------------------------------------
     gr.Examples(
         examples=[
             ["How do I create a Python function?"],
@@ -153,66 +141,44 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         inputs=msg,
         label="Example Questions"
     )
-    
-    # Event handlers
-    def user_submit(user_message, history):
-        return "", history + [[user_message, None]]
-    
-    def bot_respond(history, temperature, top_p, top_k, max_tokens):
-        user_message = history[-1][0]
-        history_context = history[:-1]
-        
-        history[-1][1] = ""
-        for response in generate_response(
-            user_message, 
-            history_context, 
-            temperature, 
-            top_p, 
-            top_k, 
-            max_tokens
-        ):
-            history[-1][1] = response
-            yield history
-    
-    # Connect events
+
+    # ------------------------------------------------------------------
+    # 6. Event wiring
+    # ------------------------------------------------------------------
     msg.submit(
-        user_submit, 
-        [msg, chatbot], 
-        [msg, chatbot], 
-        queue=False
+        user_submit, [msg, chatbot], [msg, chatbot], queue=False
     ).then(
         bot_respond,
         [chatbot, temperature, top_p, top_k, max_tokens],
-        chatbot
+        chatbot,
     )
-    
+
     submit.click(
-        user_submit,
-        [msg, chatbot],
-        [msg, chatbot],
-        queue=False
+        user_submit, [msg, chatbot], [msg, chatbot], queue=False
     ).then(
         bot_respond,
         [chatbot, temperature, top_p, top_k, max_tokens],
-        chatbot
+        chatbot,
     )
-    
+
     clear.click(lambda: None, None, chatbot, queue=False)
-    
+
     gr.Markdown(
         """
         ---
         ### 💡 Tips for Best Results:
-        - **Factual/Technical questions**: Use temperature 0.3-0.5
-        - **Creative coding solutions**: Use temperature 0.7-1.0
-        - **Code generation**: Use temperature 0.5-0.7
-        
+        - **Factual/Technical questions**: temperature 0.3-0.5
+        - **Creative coding solutions**: temperature 0.7-1.0
+        - **Code generation**: temperature 0.5-0.7
+
         ### 📝 About This Model
-        This is a fine-tuned Qwen3-4B model trained with Unsloth for efficient coding assistance.
-        
+        Fine-tuned Qwen3-4B with Unsloth for efficient coding assistance.  
         **Model**: [michsethowusu/twi_code_assistant](https://huggingface.co/michsethowusu/twi_code_assistant)
         """
     )
 
+# ------------------------------------------------------------------
+# 7. Launch
+# ------------------------------------------------------------------
 if __name__ == "__main__":
     demo.queue().launch()