Update app.py

app.py

@@ -1,3 +1,4 @@
+import gradio as gr
 import torch
 from peft import PeftModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
@@ -5,19 +6,41 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 # Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained("google/gemma-3-1b-pt")
 
-# Load base model on CPU
-base_model = AutoModelForCausalLM.from_pretrained("google/gemma-3-1b-pt")
+# Load base model on CPU with optimizations
+base_model = AutoModelForCausalLM.from_pretrained(
+    "google/gemma-3-1b-pt",
+    torch_dtype=torch.bfloat16,  # Efficient memory usage
+    low_cpu_mem_usage=True
+)
 
-# Load fine-tuned PEFT model
+# Load fine-tuned model
 model = PeftModel.from_pretrained(base_model, "hackergeek98/gemma-finetuned")
+model = model.to("cpu")  # Ensure it runs on CPU
 
-# Ensure model runs on CPU
-model = model.to("cpu")
+# Chatbot function
+def chat(message, history=[]):
+    messages = [{"role": "user", "content": message}]
+    
+    input_ids = tokenizer(message, return_tensors="pt").input_ids.to("cpu")
 
-# Test inference
-input_text = "Hello, how are you?"
-input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cpu")
+    with torch.no_grad():  # Disable gradient calculations for efficiency
+        output_ids = model.generate(input_ids, max_length=100)
+    
+    response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    
+    history.append((message, response))  # Store conversation history
+    return history, history
 
-# Generate output
-output = model.generate(input_ids, max_length=50)
-print(tokenizer.decode(output[0], skip_special_tokens=True))
+# Gradio UI
+demo = gr.ChatInterface(
+    chat,
+    chatbot=gr.Chatbot(height=400),
+    additional_inputs=[
+        gr.Textbox(value="Welcome to the chatbot!", label="System message")
+    ],
+    title="Fine-Tuned Gemma Chatbot",
+    description="This chatbot is fine-tuned on Persian text using Gemma.",
+)
+
+if __name__ == "__main__":
+    demo.launch()