Update app.py

app.py

@@ -1,50 +1,49 @@
-import os
 import torch
+from transformers import AutoProcessor, AutoModelForImageTextToText, TextStreamer
+from peft import PeftModel
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
-# Set environment variable for CUDA debugging (for local development or debugging in spaces)
-os.environ["CUDA_LAUNCH_BLOCKING"] = "1"  # Ensures errors are reported synchronously
-
-# Load the model and tokenizer
-device = "cpu"
-model = AutoModelForCausalLM.from_pretrained("adarsh3601/my_gemma3_pt", device_map="auto", torch_dtype=torch.float16).to(device)
-tokenizer = AutoTokenizer.from_pretrained("adarsh3601/my_gemma3_pt")
-
-# Function to handle input and response generation
-def generate_response(message):
-    try:
-        # Tokenize input
-        inputs = tokenizer(message, return_tensors="pt", padding=True, truncation=True, max_length=512).to(device)
-
-        # Disable gradient computation for inference
-        with torch.no_grad():
-            # Generate response from model
-            outputs = model.generate(
-                **inputs,
-                max_new_tokens=50,  # Limit the length of the output
-                temperature=0.7,
-                top_p=0.95,
-            )
-
-        # Check for NaN or Inf values in the output before decoding
-        if torch.any(torch.isnan(outputs)) or torch.any(torch.isinf(outputs)):
-            raise ValueError("Model output contains NaN or Inf")
-
-        # Decode the output
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return response
-    
-    except Exception as e:
-        print(f"Error generating response: {e}")
-        return "Error: Something went wrong."
-
-# Define the Gradio interface
-iface = gr.Interface(fn=generate_response, 
-                     inputs="text", 
-                     outputs="text", 
-                     title="Chat with Gemma3", 
-                     description="Chatbot based on Gemma3 model from Hugging Face.")
-
-# Launch the interface on Hugging Face Spaces
-iface.launch(share=True)  # To create a public link, set share=True
+
+# Load base model and processor
+base_model_id = "unsloth/gemma-3-12b-it-unsloth-bnb-4bit"
+adapter_model_id = "adarsh3601/my_gemma_pt3"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+processor = AutoProcessor.from_pretrained(base_model_id)
+model = AutoModelForImageTextToText.from_pretrained(base_model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, device_map="auto")
+
+# Apply adapter (LoRA)
+model = PeftModel.from_pretrained(model, adapter_model_id)
+model.eval()
+
+streamer = TextStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+# Helper to format messages using the chat template
+def format_chat(messages):
+    formatted = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    return formatted
+
+# Chat function
+def chat(message, history):
+    messages = []
+
+    # Format history into messages
+    for user_msg, bot_msg in history:
+        messages.append({"role": "user", "content": user_msg})
+        messages.append({"role": "assistant", "content": bot_msg})
+
+    messages.append({"role": "user", "content": message})
+    prompt = format_chat(messages)
+
+    inputs = processor(prompt, return_tensors="pt").to(device)
+
+    with torch.no_grad():
+        outputs = model.generate(**inputs, max_new_tokens=512, streamer=streamer)
+
+    decoded = processor.batch_decode(outputs, skip_special_tokens=True)[0]
+    response = decoded.split("<end_of_turn>")[0].strip().split("<start_of_turn>model")[-1].strip()
+    return response
+
+# Gradio interface
+gui = gr.ChatInterface(fn=chat, title="Gemma-3 Chatbot", description="Fine-tuned on adarsh3601/my_gemma_pt3")
+
+gui.launch()