Update app.py

app.py

@@ -1,41 +1,26 @@
-import torch
-from transformers import AutoProcessor, AutoModelForImageTextToText
-from peft import PeftModel
 import gradio as gr
+from transformers import AutoModel, AutoProcessor
+import torch
 
-# Set up device (CPU or GPU)
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Load processor and model
-model_name = "adarsh3601/my_gemma3_pt"  # Change to your model path
+# Load model and processor
+model_name = "adarsh3601/my_gemma3_pt"
+model = AutoModel.from_pretrained(model_name)
 processor = AutoProcessor.from_pretrained(model_name)
-model = AutoModelForImageTextToText.from_pretrained(model_name).to(device)
 
-# Optional: If using PEFT model with adapter
-# adapter_model_id = "your_adapter_model_id"  # Uncomment and replace if using adapter
-# model = PeftModel.from_pretrained(model, adapter_model_id)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
 
-# Define function to process the user input
-def chat(prompt):
-    # Prepare the message in the format the model expects
-    messages = [{"role": "user", "content": prompt}]
-    
-    # Process the input using the processor
+# Define a function to handle the chat interface
+def chat(messages):
     inputs = processor(messages, return_tensors="pt").to(device)
-    
-    # Generate the output from the model
-    with torch.no_grad():
-        outputs = model.generate(**inputs, max_length=200)
-
-    # Decode and return the response
+    outputs = model.generate(**inputs)
     return processor.decode(outputs[0], skip_special_tokens=True)
 
-# Gradio interface
-gr.Interface(
-    fn=chat,
-    inputs="text",
-    outputs="text",
-    title="Gemma Chat Model",
-    description="Chat with Gemma3 model",
-    live=True
-).launch(share=False)  # share=False for Hugging Face Spaces
+# Create the Gradio interface
+iface = gr.Interface(fn=chat, 
+                     inputs=gr.Textbox(label="Your Message", lines=7),
+                     outputs=gr.Textbox(label="Response"),
+                     live=True)
+
+# Launch the app
+iface.launch(debug=True)