Update app.py

app.py

@@ -1,137 +1,127 @@
 import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from peft import PeftModel, PeftConfig
-from PIL import Image
-import torchvision.datasets as datasets
-import numpy as np
-import os
+from peft import PeftModel
 
-def load_model():
-    # Create offload directory
-    os.makedirs("offload", exist_ok=True)
+# Load model and tokenizer
+def load_model(model_id):
+    # First load the base model
+    base_model_id = "microsoft/phi-2"
+    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
     
-    # Configure device map for memory efficiency
-    device_map = {
-        'base_model.model.model.embed_tokens': 0,
-        'base_model.model.model.layers.0': 0,
-        'base_model.model.model.layers.1': 0,
-        'base_model.model.model.layers.2': 0,
-        'base_model.model.model.layers.3': 0,
-        'base_model.model.model.layers.4': 'cpu',
-        'base_model.model.model.layers.5': 'cpu',
-        'base_model.model.model.layers.6': 'cpu',
-        'base_model.model.model.layers.7': 'cpu',
-        'base_model.model.model.layers.8': 'cpu',
-        'base_model.model.model.norm': 'cpu',
-        'base_model.model.lm_head': 0,
-    }
+    # Ensure tokenizer has a padding token
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
     
     base_model = AutoModelForCausalLM.from_pretrained(
-        "microsoft/Phi-3-mini-4k-instruct",
-        trust_remote_code=True,
-        device_map=device_map,  # Use custom device map
-        torch_dtype=torch.float32,
-        attn_implementation='eager',
-        offload_folder="offload"
+        base_model_id,
+        torch_dtype=torch.float16,
+        device_map="auto",
+        trust_remote_code=True
     )
     
-    model = PeftModel.from_pretrained(
-        base_model,
-        "jatingocodeo/phi-vlm",
-        device_map=device_map,
-        offload_folder="offload"
-    )
-    
-    tokenizer = AutoTokenizer.from_pretrained("jatingocodeo/phi-vlm")
-    
+    # Load and merge the LoRA adapter
+    model = PeftModel.from_pretrained(base_model, model_id)
     return model, tokenizer
 
-def generate_description(image, model, tokenizer):
-    # Convert image to RGB if needed
-    if image.mode != "RGB":
-        image = image.convert("RGB")
-    
-    # Resize image to match training size (32x32)
-    image = image.resize((32, 32))
-    
-    # Convert image to tensor and normalize
-    image_tensor = torch.FloatTensor(np.array(image)).permute(2, 0, 1) / 255.0
-    
-    # Prepare prompt with image tensor
-    prompt = f"""Below is an image. Please describe it in detail.
-
-Image: {image_tensor}
-Description: """
+def generate_response(instruction, model, tokenizer, max_length=200, temperature=0.7, top_p=0.9):
+    # Format the input text
+    input_text = instruction.strip()
     
     # Tokenize input
-    inputs = tokenizer(
-        prompt,
-        return_tensors="pt",
-        padding=True,
-        truncation=True,
-        max_length=128
-    ).to(model.device)
+    inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
     
-    # Generate description
+    # Generate response
     with torch.no_grad():
         outputs = model.generate(
-            input_ids=inputs.input_ids,
-            attention_mask=inputs.attention_mask,
-            max_new_tokens=100,
-            temperature=0.7,
-            do_sample=True,
-            top_p=0.9
+            **inputs,
+            max_new_tokens=max_length,
+            temperature=temperature,
+            top_p=top_p,
+            num_return_sequences=1,
+            pad_token_id=tokenizer.eos_token_id,
+            do_sample=True
         )
     
-    # Decode and return the generated text
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return generated_text.split("Description: ")[-1].strip()
-
-# Load model
-print("Loading model...")
-model, tokenizer = load_model()
-
-# Get CIFAR10 examples
-def get_cifar_examples():
-    cifar10_test = datasets.CIFAR10(root='./data', train=False, download=True)
-    classes = ['airplane', 'automobile', 'bird', 'cat', 'deer', 
-              'dog', 'frog', 'horse', 'ship', 'truck']
+    # Decode and return the response
+    full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
     
-    examples = []
-    used_classes = set()
+    # 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()
+
+def create_demo(model_id):
+    # Load model and tokenizer
+    model, tokenizer = load_model(model_id)
     
-    for idx in range(len(cifar10_test)):
-        img, label = cifar10_test[idx]
-        if classes[label] not in used_classes:
-            img_path = f"examples/{classes[label]}_example.jpg"
-            img.save(img_path)
-            examples.append(img_path)
-            used_classes.add(classes[label])
-            
-        if len(used_classes) == 10:
-            break
+    # Define the interface
+    def process_input(instruction, max_length, temperature, top_p):
+        try:
+            return generate_response(
+                instruction,
+                model,
+                tokenizer,
+                max_length=max_length,
+                temperature=temperature,
+                top_p=top_p
+            )
+        except Exception as e:
+            return f"Error generating response: {str(e)}"
     
-    return examples
-
-# Create Gradio interface
-def process_image(image):
-    return generate_description(image, model, tokenizer)
-
-# Get examples
-examples = get_cifar_examples()
-
-# Define interface
-iface = gr.Interface(
-    fn=process_image,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Textbox(label="Generated Description"),
-    title="Image Description Generator",
-    description="""Upload an image and get a detailed description generated by our fine-tuned VLM model.
-                  Below are sample images from CIFAR10 dataset that you can try.""",
-    examples=[[ex] for ex in examples]
-)
+    # Create the interface
+    demo = gr.Interface(
+        fn=process_input,
+        inputs=[
+            gr.Textbox(
+                label="Input Text",
+                placeholder="Enter your text here...",
+                lines=4
+            ),
+            gr.Slider(
+                minimum=50,
+                maximum=500,
+                value=150,
+                step=10,
+                label="Maximum Length"
+            ),
+            gr.Slider(
+                minimum=0.1,
+                maximum=1.0,
+                value=0.7,
+                step=0.1,
+                label="Temperature"
+            ),
+            gr.Slider(
+                minimum=0.1,
+                maximum=1.0,
+                value=0.9,
+                step=0.1,
+                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.
+        
+        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
+        """,
+        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"]
+        ]
+    )
+    return demo
 
-# Launch the interface
 if __name__ == "__main__":
-    iface.launch() 
+    # Use your model ID
+    model_id = "jatingocodeo/phi2-grpo"
+    demo = create_demo(model_id)
+    demo.launch()