Update app.py

app.py

@@ -2,31 +2,40 @@ import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from peft import PeftModel
+from PIL import Image
+import torchvision.datasets as datasets
 
-# Load model and tokenizer
 def load_model(model_id):
     # First load the base model
-    base_model_id = "microsoft/phi-2"
+    base_model_id = "microsoft/Phi-3-mini-4k-instruct"
     tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
     
-    # Ensure tokenizer has a padding token
+    # Ensure tokenizer has padding token
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
     
     base_model = AutoModelForCausalLM.from_pretrained(
         base_model_id,
-        torch_dtype=torch.float16,
+        torch_dtype=torch.float16,  # Use float16 like assignment22
         device_map="auto",
         trust_remote_code=True
     )
     
-    # Load and merge the LoRA adapter
+    # Load the LoRA adapter
     model = PeftModel.from_pretrained(base_model, model_id)
     return model, tokenizer
 
-def generate_response(instruction, model, tokenizer, max_length=200, temperature=0.7, top_p=0.9):
+def generate_description(image, model, tokenizer, max_length=100, temperature=0.7, top_p=0.9):
+    # Convert and resize image
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+    image = image.resize((32, 32))
+    
     # Format the input text
-    input_text = instruction.strip()
+    input_text = """Below is an image. Please describe it in detail.
+
+Image: [IMAGE]
+Description: """
     
     # Tokenize input
     inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
@@ -45,22 +54,31 @@ def generate_response(instruction, model, tokenizer, max_length=200, temperature
     
     # Decode and return the response
     full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    # Extract only the response part (what comes after the instruction)
-    if len(input_text) < len(full_text):
-        response = full_text[len(input_text):].strip()
-        return response
-    return full_text.strip()
+    return full_text.split("Description: ")[-1].strip()
 
 def create_demo(model_id):
     # Load model and tokenizer
     model, tokenizer = load_model(model_id)
     
-    # Define the interface
-    def process_input(instruction, max_length, temperature, top_p):
+    # Get CIFAR10 examples
+    cifar10_test = datasets.CIFAR10(root='./data', train=False, download=True)
+    examples = []
+    used_classes = set()
+    
+    for idx in range(len(cifar10_test)):
+        img, label = cifar10_test[idx]
+        class_name = cifar10_test.classes[label]
+        if class_name not in used_classes:
+            examples.append(img)
+            used_classes.add(class_name)
+        if len(used_classes) == 10:
+            break
+    
+    # Define the interface function
+    def process_image(image, max_length, temperature, top_p):
         try:
-            return generate_response(
-                instruction,
+            return generate_description(
+                image,
                 model,
                 tokenizer,
                 max_length=max_length,
@@ -68,21 +86,17 @@ def create_demo(model_id):
                 top_p=top_p
             )
         except Exception as e:
-            return f"Error generating response: {str(e)}"
+            return f"Error generating description: {str(e)}"
     
     # Create the interface
     demo = gr.Interface(
-        fn=process_input,
+        fn=process_image,
         inputs=[
-            gr.Textbox(
-                label="Input Text",
-                placeholder="Enter your text here...",
-                lines=4
-            ),
+            gr.Image(type="pil", label="Input Image"),
             gr.Slider(
                 minimum=50,
-                maximum=500,
-                value=150,
+                maximum=200,
+                value=100,
                 step=10,
                 label="Maximum Length"
             ),
@@ -101,27 +115,21 @@ def create_demo(model_id):
                 label="Top P"
             )
         ],
-        outputs=gr.Textbox(label="Completion", lines=8),
-        title="Phi-2 GRPO Model Demo",
-        description="""This is a generative model trained using GRPO (Generative Reinforcement from Preference Optimization) 
-        on the TLDR dataset. The model was trained to generate completions of around 150 characters.
+        outputs=gr.Textbox(label="Generated Description", lines=5),
+        title="Image Description Generator",
+        description="""This model generates detailed descriptions of images.
         
         You can adjust the generation parameters:
-        - **Maximum Length**: Controls the maximum length of the generated response
-        - **Temperature**: Higher values make the output more random, lower values make it more focused
-        - **Top P**: Controls the cumulative probability threshold for token sampling
+        - **Maximum Length**: Controls the length of the generated description
+        - **Temperature**: Higher values make the description more creative
+        - **Top P**: Controls the randomness in word selection
         """,
-        examples=[
-            ["The quick brown fox jumps over the lazy dog."],
-            ["In this tutorial, we will explore how to build a neural network for image classification."],
-            ["The best way to prepare for an interview is to"],
-            ["Python is a popular programming language because"]
-        ]
+        examples=[[ex] for ex in examples]
     )
     return demo
 
 if __name__ == "__main__":
     # Use your model ID
-    model_id = "jatingocodeo/phi2-grpo"
+    model_id = "jatingocodeo/phi-vlm"
     demo = create_demo(model_id)
     demo.launch()