Update app.py

app.py

@@ -1,12 +1,12 @@
 import gradio as gr
 import spaces
-from transformers import AutoProcessor, Gemma3ForConditionalGeneration
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-# Carga modelo y processor en CPU (ZeroGPU lo maneja automáticamente)
+# Carga modelo y tokenizer en CPU (ZeroGPU lo maneja automáticamente)
 model_id = "google/gemma-3-4b-it"
-processor = AutoProcessor.from_pretrained(model_id)
-model = Gemma3ForConditionalGeneration.from_pretrained(
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
     model_id,
     torch_dtype=torch.bfloat16,
     low_cpu_mem_usage=True
@@ -17,37 +17,39 @@ def generate_response(message, history):
     # Mueve el modelo a GPU temporalmente
     model.to('cuda')
     
-    # Construye mensajes para formato chat de Gemma-3 (solo texto, ya que el usuario no pidió multimodal explícitamente)
-    messages = [{"role": "system", "content": [{"type": "text", "text": "You are a helpful assistant."}]}]
+    # Construye mensajes para formato chat de Gemma-3 (solo texto)
+    messages = []
     for user_msg, ai_msg in history:
-        messages.append({"role": "user", "content": [{"type": "text", "text": user_msg}]})
+        messages.append({"role": "user", "content": user_msg})
         if ai_msg:
-            messages.append({"role": "model", "content": [{"type": "text", "text": ai_msg}]})
-    messages.append({"role": "user", "content": [{"type": "text", "text": message}]})
+            messages.append({"role": "assistant", "content": ai_msg})
+    messages.append({"role": "user", "content": message})
     
     # Aplica template de chat
-    inputs = processor.apply_chat_template(
+    prompt = tokenizer.apply_chat_template(
         messages,
-        add_generation_prompt=True,
-        tokenize=True,
-        return_dict=True,
-        return_tensors="pt"
-    ).to('cuda', dtype=torch.bfloat16)
+        tokenize=False,
+        add_generation_prompt=True
+    )
+    
+    # Tokeniza
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
     
     # Genera en GPU
-    input_len = inputs["input_ids"].shape[-1]
     with torch.no_grad():
-        generation = model.generate(
+        outputs = model.generate(
             **inputs,
             max_new_tokens=150,
             do_sample=True,
             temperature=0.7,
-            pad_token_id=processor.tokenizer.eos_token_id
+            pad_token_id=tokenizer.eos_token_id
         )
-        generation = generation[0][input_len:]
     
     # Decodifica solo tokens nuevos
-    response = processor.decode(generation, skip_special_tokens=True).strip()
+    response = tokenizer.decode(
+        outputs[0][inputs.input_ids.shape[1]:], 
+        skip_special_tokens=True
+    ).strip()
     
     # Libera GPU: mueve modelo de vuelta a CPU y limpia caché
     model.to('cpu')
@@ -55,20 +57,17 @@ def generate_response(message, history):
     
     return response
 
-# Interfaz de chat con Gradio (solo texto, para simplicidad)
+# Interfaz de chat con Gradio (solo texto)
 demo = gr.ChatInterface(
     fn=generate_response,
-    title="Chat Interface para Hugging Face ZeroGPU Space con Gemma-3-4B-IT",
-    description="Un chatbot usando google/gemma-3-4b-it. Optimizado para ZeroGPU: GPU se asigna dinámicamente solo durante la generación con @spaces.GPU. (Solo texto; ignora multimodal por ahora).",
+    title="Gemma-3-4B-IT Chat en ZeroGPU",
+    description="Chatbot con google/gemma-3-4b-it. GPU dinámica via @spaces.GPU.",
     examples=[
         ["Hola, ¿cómo estás?"],
-        ["Explica qué es la IA."],
+        ["Explica qué es la IA en 3 líneas."],
         ["Cuéntame un chiste."]
     ],
-    cache_examples=True,
-    retry_btn="Reintentar",
-    undo_btn="Deshacer",
-    clear_btn="Limpiar chat"
+    cache_examples=True
 )
 
 if __name__ == "__main__":