Update app.py

app.py

@@ -1,13 +1,22 @@
 from vllm import LLM
 from vllm.sampling_params import SamplingParams
-import mergekit
 from fastapi import FastAPI
 from uvicorn import run
 
 app = FastAPI()
 
 model_name = "Hjgugugjhuhjggg/mergekit-ties-tzamfyy"
-mergekit_model = mergekit.MergeKitModel(model_name)
+
+# Verifica si mergekit tiene el atributo correcto
+try:
+    import mergekit
+    if hasattr(mergekit, "MergeKitModel"):
+        mergekit_model = mergekit.MergeKitModel(model_name)
+    else:
+        raise ImportError("MergeKitModel no encontrado en mergekit. Verifica la instalación.")
+except ImportError as e:
+    print(f"Error importando mergekit: {e}")
+    mergekit_model = None
 
 SYSTEM_PROMPT = "¡Hola! Soy un modelo de lenguaje avanzado. Estoy aquí para ayudarte con cualquier pregunta o tema que desees discutir. ¿En qué puedo ayudarte hoy?"
 
@@ -32,7 +41,6 @@ def divide_texto(texto, max_tokens=512):
 async def chat(texto: str):
     mensajes = divide_texto(texto)
 
-    # Nuevo mensaje
     mensaje = {"role": "user", "content": []}
 
     for mensaje_texto in mensajes:
@@ -45,11 +53,18 @@ async def chat(texto: str):
 
     sampling_params = SamplingParams(max_tokens=512)
 
-    llm = LLM(model=mergekit_model, config_format="llama", load_format="llama", tokenizer_mode="llama", tensor_parallel_size=8)
-
-    outputs = llm.chat(messages, sampling_params=sampling_params)
-
-    return {"response": outputs[0].outputs[0].text}
+    try:
+        llm = LLM(
+            model=mergekit_model if mergekit_model else model_name,
+            config_format="llama",
+            load_format="llama",
+            tokenizer_mode="llama",
+            tensor_parallel_size=8
+        )
+        outputs = llm.chat(messages, sampling_params=sampling_params)
+        return {"response": outputs[0].outputs[0].text}
+    except Exception as e:
+        return {"error": str(e)}
 
 if __name__ == "__main__":
     run(app, host="0.0.0.0", port=8000)