Remove LoRA dependencies, use base Gemma-2-9B model

app.py

@@ -1,38 +1,48 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from peft import PeftModel
 import torch
 import gradio as gr
 
 # Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b-it")
 
-# Load base model on CPU
-base_model = AutoModelForCausalLM.from_pretrained(
+# Load base model directly (no LoRA adapters)
+model = AutoModelForCausalLM.from_pretrained(
     "google/gemma-2-9b-it",
     torch_dtype=torch.bfloat16,
-    device_map="cpu",
+    device_map="auto",
     low_cpu_mem_usage=True
 )
 
-# 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)
+    # Format the prompt for chat
+    formatted_prompt = f"<start_of_turn>user\n{prompt}<end_of_turn>\n<start_of_turn>model\n"
+    
+    inputs = tokenizer(formatted_prompt, return_tensors="pt")
+    
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9,
+            pad_token_id=tokenizer.eos_token_id
+        )
+    
+    # Decode only the generated part (excluding the input)
+    generated_text = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
+    return generated_text.strip()
 
 # 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 Gemma2-9B",
-    description="Chat with a flirty AI powered by Gemma-2-9B!"
+    inputs=[
+        gr.Textbox(label="Your message", placeholder="Type your message here..."),
+        gr.Slider(minimum=10, maximum=200, value=50, label="Max New Tokens")
+    ],
+    outputs=gr.Textbox(label="AI Response"),
+    title="Flirt-AI Gemma2-9B (Base Model)",
+    description="Chat with AI powered by the base Gemma-2-9B model!"
 )
 
 if __name__ == "__main__":