Update app.py

app.py

@@ -2,14 +2,11 @@ import os
 import torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
-from peft import PeftModel
 
 # =========================
 # CONFIG
 # =========================
-BASE_MODEL = "google/gemma-3-270m-it"
-LORA_MODEL = "AxionLab-official/DogeAI-v1.0-instruct"
-
+MODEL_ID = "google/gemma-2-2b-it"
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
 # =========================
@@ -17,31 +14,21 @@ HF_TOKEN = os.environ.get("HF_TOKEN")
 # =========================
 print("🔄 Loading tokenizer...")
 tokenizer = AutoTokenizer.from_pretrained(
-    BASE_MODEL,
+    MODEL_ID,
     token=HF_TOKEN,
     use_fast=False
 )
 
 # =========================
-# LOAD BASE MODEL (CPU SAFE)
+# LOAD MODEL
 # =========================
-print("🔄 Loading base model...")
+print("🔄 Loading model...")
 model = AutoModelForCausalLM.from_pretrained(
-    BASE_MODEL,
+    MODEL_ID,
     token=HF_TOKEN,
     torch_dtype=torch.float32,
     device_map="auto"
 )
-
-# =========================
-# APPLY LORA
-# =========================
-print("🔄 Applying LoRA...")
-model = PeftModel.from_pretrained(
-    model,
-    LORA_MODEL,
-    device_map="auto"
-)
 model.eval()
 
 # =========================
@@ -62,7 +49,7 @@ def chat(user_input, system_prompt, temperature, top_p, max_tokens):
     inputs = tokenizer(
         prompt,
         return_tensors="pt"
-    )
+    ).to(model.device)
 
     with torch.no_grad():
         output = model.generate(
@@ -81,7 +68,6 @@ def chat(user_input, system_prompt, temperature, top_p, max_tokens):
         skip_special_tokens=True
     )
 
-    # Remove prompt da resposta
     return decoded[len(prompt):].strip()
 
 # =========================
@@ -92,7 +78,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         """
         # 🐕 DogeAI v1.0  
         Modelo **experimental**, pequeno e focado em respostas claras.  
-
         ⚠️ Pode errar ou alucinar  
         ❌ Não use como fonte confiável  
         ✅ Use para estudo e experimentação