Update app.py

app.py

@@ -1,98 +1,51 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from peft import PeftModel
-import torch
+from transformers import pipeline
 import os
 
-HF_TOKEN = os.environ.get("HF_TOKEN")
+HF_TOKEN = os.environ["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 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()
+# Initialize the text generation pipelines
+pipe = pipeline("text-generation", model="akhaliq/MyGemmaGradioCoder")
+pipe2 = pipeline("text-generation", model="google/gemma-3-270m-it", token=HF_TOKEN)
 
-# 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"):
+def generate_code(user_input, model_choice="Model 1"):
     """
-    Generate code based on user input using the selected model
+    Generate code based on user input using the selected Gemma model
     """
-    # 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()
+    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 ""
     else:
-        response = generated_text
-    
-    # Clean up any remaining turn markers
-    response = response.replace("<end_of_turn>", "").strip()
-    
-    return response
+        # 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
 
 def generate_both(user_input):
     """
     Generate code from both models for comparison
     """
-    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)}"
-    
+    output1 = generate_code(user_input, "Model 1 (MyGemmaGradioCoder)")
+    output2 = generate_code(user_input, "Model 2 (gemma-3-270m-it)")
     return output1, output2
 
 # Create the Gradio interface
@@ -102,8 +55,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**: Gemma with Gradio Coder Adapter (Fine-tuned)
-        - **Model 2**: Base Gemma Model
+        - **Model 1**: akhaliq/MyGemmaGradioCoder
+        - **Model 2**: google/gemma-3-270m-it
         
         Simply describe what you want to build, and see how each model responds!
         """
@@ -131,7 +84,6 @@ 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"
@@ -141,27 +93,31 @@ 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: With Gradio Coder Adapter")
+                    gr.Markdown("### Model 1: MyGemmaGradioCoder")
                     output_code1 = gr.Code(
                         label="Generated Code (Model 1)",
                         language="python",
                         lines=15,
                         interactive=True,
-                        show_label=False
+                        show_line_numbers=True,
+                        wrap_lines=True,
+                        autocomplete=True
                     )
-                    copy_btn1 = gr.Button("📋 Copy Code", size="sm")
+                    copy_btn1 = gr.Button("📋 Copy Code", scale=1)
                 
                 with gr.Column():
-                    gr.Markdown("### Model 2: Base Gemma Model")
+                    gr.Markdown("### Model 2: gemma-3-270m-it")
                     output_code2 = gr.Code(
                         label="Generated Code (Model 2)",
                         language="python",
                         lines=15,
                         interactive=True,
-                        show_label=False
+                        show_line_numbers=True,
+                        wrap_lines=True,
+                        autocomplete=True
                     )
-                    copy_btn2 = gr.Button("📋 Copy Code", size="sm")
-    
+                    copy_btn2 = gr.Button("📋 Copy Code", scale=1)
+            
     # Add event handlers
     generate_btn.click(
         fn=generate_both,
@@ -178,33 +134,26 @@ with gr.Blocks(title="Text to Code Generator - Model Comparison", theme=gr.theme
     
     # Copy functionality for both outputs
     copy_btn1.click(
-        None,
+        fn=None,
         inputs=output_code1,
         outputs=None,
         js="""
         (code) => {
             navigator.clipboard.writeText(code);
-            const btn = document.querySelector('button:has-text("📋 Copy Code")');
-            const originalText = btn.textContent;
-            btn.textContent = '✓ Copied!';
-            setTimeout(() => btn.textContent = originalText, 2000);
+            alert('Code from Model 1 copied to clipboard!');
             return null;
         }
         """
     )
     
     copy_btn2.click(
-        None,
+        fn=None,
         inputs=output_code2,
         outputs=None,
         js="""
         (code) => {
             navigator.clipboard.writeText(code);
-            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);
+            alert('Code from Model 2 copied to clipboard!');
             return null;
         }
         """
@@ -219,7 +168,9 @@ 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
         
-        **Note:** The adapter model is specifically fine-tuned for generating Gradio code!
+        **Models:**
+        - [akhaliq/MyGemmaGradioCoder](https://huggingface.co/akhaliq/MyGemmaGradioCoder)
+        - [google/gemma-3-270m-it](https://huggingface.co/google/gemma-3-270m-it)
         """
     )