Update app.py

app.py

@@ -212,49 +212,46 @@ EXAMPLES = [
 # Definir la función simple_chat
 @spaces.GPU()
 def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096, top_p: float = 1, top_k: int = 10, penalty: float = 1.0):
-    # Cargar el modelo preentrenado
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID,
-        torch_dtype=torch.bfloat16,
-        low_cpu_mem_usage=True,
-        trust_remote_code=True
-    )
-
-    conversation = []
-
-    if "file" in message and message["file"]:
-        file_path = message["file"]
-        choice, contents = mode_load(file_path)
-        if choice == "image":
-            conversation.append({"role": "user", "image": contents, "content": message["text"]})
-        elif choice == "doc":
-            format_msg = contents + "\n\n\n" + "{} files uploaded.\n" + message["text"]
-            conversation.append({"role": "user", "content": format_msg})
-    else:
-        conversation.append({"role": "user", "content": message["text"]})
-
-    # Preparar entrada para el modelo
-    input_ids = tokenizer.apply_chat_template(conversation, tokenize=True, add_generation_prompt=True,
-                                              return_tensors="pt", return_dict=True).to(model.device)
-    
-    # Configurar parámetros de generación
-    generate_kwargs = dict(
-        max_length=max_length,
-        do_sample=True,
-        top_p=top_p,
-        top_k=top_k,
-        temperature=temperature,
-        repetition_penalty=penalty,
-        eos_token_id=[151329, 151336, 151338],
-    )
-
-    # Generar respuesta
-    with torch.no_grad():
-        generated_ids = model.generate(input_ids['input_ids'], **generate_kwargs)
-        generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+    try:
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_ID,
+            torch_dtype=torch.bfloat16,
+            low_cpu_mem_usage=True,
+            trust_remote_code=True
+        )
+
+        conversation = []
+
+        if "file" in message and message["file"]:
+            file_path = message["file"]
+            choice, contents = mode_load(file_path)
+            if choice == "image":
+                conversation.append({"role": "user", "image": contents, "content": message["text"]})
+            elif choice == "doc":
+                format_msg = contents + "\n\n\n" + "{} files uploaded.\n" + message["text"]
+                conversation.append({"role": "user", "content": format_msg})
+        else:
+            conversation.append({"role": "user", "content": message["text"]})
 
-    # Devolver la respuesta completa
-    return PlainTextResponse(generated_text)
+        input_ids = tokenizer.apply_chat_template(conversation, tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(model.device)
+        
+        generate_kwargs = dict(
+            max_length=max_length,
+            do_sample=True,
+            top_p=top_p,
+            top_k=top_k,
+            temperature=temperature,
+            repetition_penalty=penalty,
+            eos_token_id=[151329, 151336, 151338],
+        )
+
+        with torch.no_grad():
+            generated_ids = model.generate(input_ids['input_ids'], **generate_kwargs)
+            generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+
+        return PlainTextResponse(generated_text)
+    except Exception as e:
+        return PlainTextResponse(f"Error: {str(e)}")
 
 @app.post("/chat/")
 async def test_endpoint(message: dict):