comparing base with finetuned

main.py

@@ -22,8 +22,8 @@ MODEL = "meta-llama/Llama-3.2-1B-Instruct"
 
 app = FastAPI()
 
-# base model and tokenizer
-base_model = AutoModelForCausalLM.from_pretrained(
+# base model for finetuned (LoRA) inference
+finetuned_base = AutoModelForCausalLM.from_pretrained(
     MODEL,
     token=HF_TOKEN,
     dtype=torch.bfloat16,  # faster than float32, matches GPU training
@@ -31,8 +31,19 @@ base_model = AutoModelForCausalLM.from_pretrained(
     low_cpu_mem_usage=True,
     attn_implementation="sdpa",  # PyTorch optimized attention
 )
+finetuned_base.config.use_cache = True
 
+# separate base model for comparison (no LoRA)
+base_model = AutoModelForCausalLM.from_pretrained(
+    MODEL,
+    token=HF_TOKEN,
+    dtype=torch.bfloat16,
+    device_map="cpu",
+    low_cpu_mem_usage=True,
+    attn_implementation="sdpa",
+)
 base_model.config.use_cache = True
+base_model.eval()
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL, token=HF_TOKEN)
 
@@ -50,7 +61,7 @@ languages = list(adapter_paths.keys())
 
 # Create PeftModel with first adapter
 peft_model = PeftModel.from_pretrained(
-    base_model,
+    finetuned_base,
     adapter_paths[languages[0]],
     adapter_name=languages[0],
     is_trainable=False
@@ -63,6 +74,44 @@ for lang in languages[1:]:
 peft_model.eval()
 print("All adapters ready.")
 
+# base model generation (no LoRA)
+def generate_base_model_stream(messages, max_length=256, temperature=0.7):
+
+    print(f"Base model (no LoRA)")
+    print(f"Messages: {messages}")
+
+    inputs = tokenizer.apply_chat_template(
+        messages,
+        tokenize=True,
+        add_generation_prompt=True,
+        return_tensors="pt",
+        return_dict=True
+    ).to(base_model.device)
+
+    streamer = TextIteratorStreamer(
+        tokenizer,
+        skip_prompt=True,
+        skip_special_tokens=True)
+
+    generation_kwargs = {
+        **inputs,
+        "max_new_tokens": max_length,
+        "temperature": temperature,
+        "do_sample": True,
+        "pad_token_id": tokenizer.eos_token_id,
+        "streamer": streamer,
+        "num_beams": 1,
+        "use_cache": True,
+    }
+
+    thread = Thread(target=base_model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    for text in streamer:
+        yield text
+
+    thread.join()
+
 # the input will be a list of messages that include system, user, and assistant prompts
 def generate_text_stream(messages, language, max_length=256, temperature=0.7):
 
@@ -249,14 +298,32 @@ async def chat_completions(request: ChatCompletionRequest):
         }
     )
 
-def chat_gradio(message, history, language, system_prompt, max_length, temperature):
+def chat_base_model(message, history, system_prompt, max_length, temperature):
+    messages = []
 
+    if system_prompt:
+        messages.append({"role": "system", "content": system_prompt})
+
+    messages.extend(history[-10:])
+
+    user_msg = {"role": "user", "content": message}
+    messages.append(user_msg)
+
+    assistant_msg = {"role": "assistant", "content": ""}
+    for token in generate_base_model_stream(
+        messages,
+        max_length,
+        temperature
+    ):
+        assistant_msg["content"] += token
+        yield history + [user_msg, assistant_msg]
+
+def chat_finetuned(message, history, language, system_prompt, max_length, temperature):
     messages = []
 
     if system_prompt:
         messages.append({"role": "system", "content": system_prompt})
 
-    # only uses the last 10 messages to keep within context limit
     messages.extend(history[-10:])
 
     user_msg = {"role": "user", "content": message}
@@ -273,48 +340,57 @@ def chat_gradio(message, history, language, system_prompt, max_length, temperatu
         yield history + [user_msg, assistant_msg]
 
 with gr.Blocks(
-    title="Language Learning Chatbot", 
+    title="Language Learning Chatbot",
     theme=gr.themes.Soft()
 ) as demo:
 
     with gr.Row():
-        with gr.Column(scale=2):
-            chatbot = gr.Chatbot(
-                label="Conversation",
-                height=500,
+        with gr.Column(scale=1):
+            gr.Markdown("### Base Model (No LoRA)")
+            chatbot_base = gr.Chatbot(
+                label="Base Model",
+                height=400,
                 type="messages"
             )
 
-            # User input
-            with gr.Row():
-                msg = gr.Textbox(
-                    label="Your message",
-                    placeholder="Type your message here and press Enter...",
-                    lines=2,
-                    scale=4
-                )
+        with gr.Column(scale=1):
+            gr.Markdown("### Finetuned Model (LoRA)")
+            chatbot_finetuned = gr.Chatbot(
+                label="Finetuned Model",
+                height=400,
+                type="messages"
+            )
 
-            with gr.Row():
-                submit_btn = gr.Button("Send", variant="primary", scale=1)
-                clear_btn = gr.Button("Clear Chat", scale=1)
+    with gr.Row():
+        msg = gr.Textbox(
+            label="Your message",
+            placeholder="Type your message here and press Enter...",
+            lines=2,
+            scale=4
+        )
+
+    with gr.Row():
+        submit_btn = gr.Button("Send", variant="primary", scale=1)
+        clear_btn = gr.Button("Clear Both Chats", scale=1)
+
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("### Settings")
 
-        with gr.Column(scale=1):
-            gr.Markdown("### ⚙️ Settings")
-            
             language_dropdown = gr.Dropdown(
                 choices=list(adapter_paths.keys()),
-                label="Language",
+                label="Language (for Finetuned Model)",
                 value=list(adapter_paths.keys())[0],
-                info="Select the language model to use"
+                info="Select the language adapter to use"
             )
 
             system_prompt_input = gr.Textbox(
                 label="System Prompt (Optional)",
                 placeholder="e.g., You are a helpful assistant...",
                 lines=3,
-                info="Set the assistant's behavior"
+                info="Shared between both models"
             )
-            
+
             max_length_slider = gr.Slider(
                 minimum=50,
                 maximum=512,
@@ -333,31 +409,41 @@ with gr.Blocks(
                 info="Higher = more creative"
             )
 
-    # handling enter key in textbox
+    # handling enter key in textbox - send to both models
+    msg.submit(
+        fn=chat_base_model,
+        inputs=[msg, chatbot_base, system_prompt_input, max_length_slider, temperature_slider],
+        outputs=[chatbot_base],
+    )
     msg.submit(
-        fn=chat_gradio,
-        inputs=[msg, chatbot, language_dropdown, system_prompt_input, max_length_slider, temperature_slider],
-        outputs=[chatbot],  # Update chatbot with streaming response
+        fn=chat_finetuned,
+        inputs=[msg, chatbot_finetuned, language_dropdown, system_prompt_input, max_length_slider, temperature_slider],
+        outputs=[chatbot_finetuned],
     ).then(
-        fn=lambda: gr.update(value=""),  # Clear input after sending
+        fn=lambda: gr.update(value=""),
         outputs=[msg]
     )
 
-    # Handle button click
+    # Handle button click - send to both models
     submit_btn.click(
-        fn=chat_gradio,
-        inputs=[msg, chatbot, language_dropdown, system_prompt_input, max_length_slider, temperature_slider],
-        outputs=[chatbot],
+        fn=chat_base_model,
+        inputs=[msg, chatbot_base, system_prompt_input, max_length_slider, temperature_slider],
+        outputs=[chatbot_base],
+    )
+    submit_btn.click(
+        fn=chat_finetuned,
+        inputs=[msg, chatbot_finetuned, language_dropdown, system_prompt_input, max_length_slider, temperature_slider],
+        outputs=[chatbot_finetuned],
     ).then(
         fn=lambda: gr.update(value=""),
         outputs=[msg]
     )
-    
-    # Clear chat button
+
+    # Clear both chats
     clear_btn.click(
-        fn=lambda: None,  # Return None to clear chatbot
-        outputs=[chatbot],
-        queue=False  # Don't queue this action
+        fn=lambda: (None, None),
+        outputs=[chatbot_base, chatbot_finetuned],
+        queue=False
     )
 
 demo.queue(False)