Update app.py

app.py

@@ -1,31 +1,21 @@
-import os
 import spaces
 import gradio as gr
-from huggingface_hub import hf_hub_download
-from llama_cpp import Llama
-
-REPO_ID = "baa-ai/Kimi-K2.6-RAM-GGUF"
-# Nota: Si el modelo está shardeado en varias partes (-00001-of-00014),
-# llama.cpp requiere que descargues la primera parte y automáticamente detecta el resto,
-# o debes usar un archivo único cuantizado menor (p. ej. Q2_K o un modelo de menor tamaño).
-FILENAME = "baa-ai/Kimi-K2.6-RAM-GGUF"
-
-# Descargar archivo del Hub
-model_path = hf_hub_download(
-    repo_id=REPO_ID,
-    filename=FILENAME
-)
-
-# Inicializar motor GGUF con soporte GPU
-llm = Llama(
-    model_path=model_path,
-    n_gpu_layers=-1,       # Enviar capas a la GPU
-    n_ctx=4096,
-    verbose=False
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from threading import Thread
+
+MODEL_ID = "deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct"
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    trust_remote_code=True
 )
 
 @spaces.GPU(duration=120)
-def generate_response(message, history, system_prompt=""):
+def generate(message, history, system_prompt=""):
     messages = []
     if system_prompt:
         messages.append({"role": "system", "content": system_prompt})
@@ -36,27 +26,36 @@ def generate_response(message, history, system_prompt=""):
             messages.append({"role": "assistant", "content": bot_msg})
             
     messages.append({"role": "user", "content": message})
-
-    response_stream = llm.create_chat_completion(
-        messages=messages,
-        max_tokens=2048,
-        temperature=0.6,
-        top_p=0.9,
-        stream=True
+    
+    inputs = tokenizer.apply_chat_template(
+        messages, 
+        tokenize=True, 
+        add_generation_prompt=True, 
+        return_tensors="pt"
+    ).to(model.device)
+
+    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
+    generate_kwargs = dict(
+        inputs=inputs,
+        streamer=streamer,
+        max_new_tokens=2048,
+        temperature=0.3,
+        top_p=0.95
     )
 
+    thread = Thread(target=model.generate, kwargs=generate_kwargs)
+    thread.start()
+
     partial_text = ""
-    for chunk in response_stream:
-        delta = chunk["choices"][0]["delta"]
-        if "content" in delta:
-            partial_text += delta["content"]
-            yield partial_text
+    for new_token in streamer:
+        partial_text += new_token
+        yield partial_text
 
 demo = gr.ChatInterface(
-    fn=generate_response,
-    title="Kimi-K2.6 Inference",
+    fn=generate,
+    title="DeepSeek Coder V2 Lite (ZeroGPU)",
     additional_inputs=[
-        gr.Textbox("Eres un asistente experto en programación.", label="System Prompt")
+        gr.Textbox("Eres un asistente experto en programación y resolución de problemas.", label="System Prompt")
     ]
 )