Update app.py

app.py

@@ -1,51 +1,98 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
+import torch
 import os
 
-HF_TOKEN = os.environ["HF_TOKEN"]
+HF_TOKEN = os.environ.get("HF_TOKEN")
 
+# Load the base model and adapter for Model 1
+base_model_name = "google/gemma-2b-it"  # or the correct base model
+adapter_model_name = "akhaliq/gemma-3-270m-gradio-coder-adapter"
 
-# Initialize the text generation pipelines
-pipe = pipeline("text-generation", model="akhaliq/gemma-3-270m-gradio-coder-adapter")
-pipe2 = pipeline("text-generation", model="google/gemma-3-270m-it", token=HF_TOKEN)
+# Initialize Model 1 (with adapter)
+print("Loading Model 1 with adapter...")
+tokenizer1 = AutoTokenizer.from_pretrained(adapter_model_name)
+base_model1 = AutoModelForCausalLM.from_pretrained(
+    base_model_name,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    device_map="auto" if torch.cuda.is_available() else None,
+    token=HF_TOKEN
+)
+model1 = PeftModel.from_pretrained(base_model1, adapter_model_name)
+model1.eval()
 
-def generate_code(user_input, model_choice="Model 1"):
+# Initialize Model 2 (standard model)
+print("Loading Model 2...")
+model2_name = "google/gemma-2b-it"  # Using gemma-2b-it as gemma-3-270m-it might not exist
+tokenizer2 = AutoTokenizer.from_pretrained(model2_name, token=HF_TOKEN)
+model2 = AutoModelForCausalLM.from_pretrained(
+    model2_name,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    device_map="auto" if torch.cuda.is_available() else None,
+    token=HF_TOKEN
+)
+model2.eval()
+
+def generate_code(user_input, model, tokenizer, model_name="Model"):
     """
-    Generate code based on user input using the selected Gemma model
+    Generate code based on user input using the selected model
     """
-    messages = [
-        {"role": "user", "content": user_input},
-    ]
-    
-    # Select pipeline based on model choice
-    selected_pipe = pipe if model_choice == "Model 1 (MyGemmaGradioCoder)" else pipe2
-    
-    # Generate response from the model
-    response = selected_pipe(messages, max_new_tokens=512, temperature=0.7, do_sample=True)
-    
-    # Extract the generated text from the response
-    generated_text = response[0]['generated_text']
-    
-    # If the response contains the full conversation, extract just the assistant's response
-    if isinstance(generated_text, list):
-        # Handle conversation format
-        for msg in generated_text:
-            if msg.get('role') == 'assistant':
-                return msg.get('content', '')
-        # If no assistant message found, return the last message content
-        return generated_text[-1].get('content', '') if generated_text else ""
+    # Format the prompt for code generation
+    prompt = f"<start_of_turn>user\n{user_input}<end_of_turn>\n<start_of_turn>model\n"
+    
+    # Tokenize input
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
+    
+    # Move to same device as model
+    if torch.cuda.is_available():
+        inputs = {k: v.cuda() for k, v in inputs.items()}
+    
+    # Generate response
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=512,
+            temperature=0.7,
+            do_sample=True,
+            top_p=0.9,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+        )
+    
+    # Decode the output
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # Extract only the model's response
+    if "<start_of_turn>model" in generated_text:
+        response = generated_text.split("<start_of_turn>model")[-1].strip()
+    elif user_input in generated_text:
+        response = generated_text.split(user_input)[-1].strip()
     else:
-        # Handle string format - try to extract the code after the user input
-        if user_input in generated_text:
-            return generated_text.split(user_input)[-1].strip()
-        return generated_text
+        response = generated_text
+    
+    # Clean up any remaining turn markers
+    response = response.replace("<end_of_turn>", "").strip()
+    
+    return response
 
 def generate_both(user_input):
     """
     Generate code from both models for comparison
     """
-    output1 = generate_code(user_input, "Model 1 (MyGemmaGradioCoder)")
-    output2 = generate_code(user_input, "Model 2 (gemma-3-270m-it)")
+    if not user_input.strip():
+        return "", ""
+    
+    try:
+        output1 = generate_code(user_input, model1, tokenizer1, "Model 1 (Adapter)")
+    except Exception as e:
+        output1 = f"Error with Model 1: {str(e)}"
+    
+    try:
+        output2 = generate_code(user_input, model2, tokenizer2, "Model 2 (Base)")
+    except Exception as e:
+        output2 = f"Error with Model 2: {str(e)}"
+    
     return output1, output2
 
 # Create the Gradio interface
@@ -55,8 +102,8 @@ with gr.Blocks(title="Text to Code Generator - Model Comparison", theme=gr.theme
         # 🚀 Text to Code Generator - Model Comparison
         
         Compare code generation from two different Gemma models:
-        - **Model 1**: akhaliq/MyGemmaGradioCoder
-        - **Model 2**: google/gemma-3-270m-it
+        - **Model 1**: Gemma with Gradio Coder Adapter (Fine-tuned)
+        - **Model 2**: Base Gemma Model
         
         Simply describe what you want to build, and see how each model responds!
         """
@@ -84,6 +131,7 @@ with gr.Blocks(title="Text to Code Generator - Model Comparison", theme=gr.theme
                     ["Create a React component for a todo list item"],
                     ["Write a SQL query to find the top 5 customers by total purchase amount"],
                     ["Create a Python class for a bank account with deposit and withdraw methods"],
+                    ["Build a simple Gradio interface for text summarization"],
                 ],
                 inputs=input_text,
                 label="Example Prompts"
@@ -93,31 +141,27 @@ with gr.Blocks(title="Text to Code Generator - Model Comparison", theme=gr.theme
             # Output section - Two columns for comparison
             with gr.Row():
                 with gr.Column():
-                    gr.Markdown("### Model 1: MyGemmaGradioCoder")
+                    gr.Markdown("### Model 1: With Gradio Coder Adapter")
                     output_code1 = gr.Code(
                         label="Generated Code (Model 1)",
                         language="python",
                         lines=15,
                         interactive=True,
-                        show_line_numbers=True,
-                        wrap_lines=True,
-                        autocomplete=True
+                        show_label=False
                     )
-                    copy_btn1 = gr.Button("📋 Copy Code", scale=1)
+                    copy_btn1 = gr.Button("📋 Copy Code", size="sm")
                 
                 with gr.Column():
-                    gr.Markdown("### Model 2: gemma-3-270m-it")
+                    gr.Markdown("### Model 2: Base Gemma Model")
                     output_code2 = gr.Code(
                         label="Generated Code (Model 2)",
                         language="python",
                         lines=15,
                         interactive=True,
-                        show_line_numbers=True,
-                        wrap_lines=True,
-                        autocomplete=True
+                        show_label=False
                     )
-                    copy_btn2 = gr.Button("📋 Copy Code", scale=1)
-            
+                    copy_btn2 = gr.Button("📋 Copy Code", size="sm")
+    
     # Add event handlers
     generate_btn.click(
         fn=generate_both,
@@ -134,26 +178,33 @@ with gr.Blocks(title="Text to Code Generator - Model Comparison", theme=gr.theme
     
     # Copy functionality for both outputs
     copy_btn1.click(
-        fn=None,
+        None,
         inputs=output_code1,
         outputs=None,
         js="""
         (code) => {
             navigator.clipboard.writeText(code);
-            alert('Code from Model 1 copied to clipboard!');
+            const btn = document.querySelector('button:has-text("📋 Copy Code")');
+            const originalText = btn.textContent;
+            btn.textContent = '✓ Copied!';
+            setTimeout(() => btn.textContent = originalText, 2000);
             return null;
         }
         """
     )
     
     copy_btn2.click(
-        fn=None,
+        None,
         inputs=output_code2,
         outputs=None,
         js="""
         (code) => {
             navigator.clipboard.writeText(code);
-            alert('Code from Model 2 copied to clipboard!');
+            const btns = document.querySelectorAll('button:has-text("📋 Copy Code")');
+            const btn = btns[1];
+            const originalText = btn.textContent;
+            btn.textContent = '✓ Copied!';
+            setTimeout(() => btn.textContent = originalText, 2000);
             return null;
         }
         """
@@ -168,9 +219,7 @@ with gr.Blocks(title="Text to Code Generator - Model Comparison", theme=gr.theme
         - Include details about inputs, outputs, and edge cases
         - You can edit the generated code directly in the output box
         
-        **Models:**
-        - [akhaliq/MyGemmaGradioCoder](https://huggingface.co/akhaliq/MyGemmaGradioCoder)
-        - [google/gemma-3-270m-it](https://huggingface.co/google/gemma-3-270m-it)
+        **Note:** The adapter model is specifically fine-tuned for generating Gradio code!
         """
     )