Update app.py

app.py

@@ -1,25 +1,39 @@
-import gradio as gr
-from unsloth import FastLanguageModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
 import torch
+import gradio as gr
 
-model, tokenizer = FastLanguageModel.from_pretrained(
-    "lingadevaruhp/flirt-ai-gemma2-2b",
-    max_seq_length=2048,
-    dtype=torch.float16,
-    load_in_4bit=True,
-    device_map="auto"
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b-it")
+
+# Load base model on CPU
+base_model = AutoModelForCausalLM.from_pretrained(
+    "google/gemma-2-9b-it",
+    torch_dtype=torch.bfloat16,
+    device_map="cpu",
+    low_cpu_mem_usage=True
 )
-FastLanguageModel.for_inference(model)
 
-def chat(prompt):
-    inputs = tokenizer(f"<s>### Instruction:\n{prompt}\n### Response:\n", return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
-    outputs = model.generate(**inputs, max_new_tokens=100)
+# Load LoRA adapters (replace with your repo once pushed)
+model = PeftModel.from_pretrained(
+    base_model,
+    "lingadevaruhp/flirt-ai-gemma2-9b",  # Update after pushing
+    device_map="cpu"
+)
+
+def generate_response(prompt, max_new_tokens=50):
+    inputs = tokenizer(prompt, return_tensors="pt")
+    outputs = model.generate(**inputs, max_new_tokens=max_new_tokens)
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-gr.Interface(
-    fn=chat,
-    inputs="text",
+# Gradio interface
+iface = gr.Interface(
+    fn=generate_response,
+    inputs=["text", gr.Slider(minimum=10, maximum=200, value=50, label="Max New Tokens")],
     outputs="text",
-    title="Flirt.AI: Kannada Flirty Chatbot",
-    description="Chat with a flirty AI using Kannada slang and English! 😎"
-).launch()
+    title="Flirt-AI Gemma2-9B",
+    description="Chat with a flirty AI powered by Gemma-2-9B!"
+)
+
+if __name__ == "__main__":
+    iface.launch()