Update app.py

app.py

@@ -27,53 +27,6 @@ async def test_endpoint(message: dict):
     return response
 
 
-@app.post("/chat/")
-async def chat_endpoint(message: dict):
-    if "text" not in message:
-        raise HTTPException(status_code=400, detail="Missing 'text' in request body")
-    
-    chat_message = message["text"]
-    response_text = generate_chat_response(chat_message)
-    
-    return {"response": response_text}
-
-def generate_chat_response(text: str):
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID,
-        torch_dtype=torch.bfloat16,
-        low_cpu_mem_usage=True,
-        trust_remote_code=True
-    )
-    
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
-
-    conversation = [{"role": "user", "content": text}]
-    input_ids = tokenizer.apply_chat_template(conversation, tokenize=True, add_generation_prompt=True,
-                                              return_tensors="pt", return_dict=True).to(model.device)
-    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
-
-    generate_kwargs = dict(
-        max_length=4096,
-        streamer=streamer,
-        do_sample=True,
-        top_p=0.9,
-        top_k=50,
-        temperature=0.7,
-        repetition_penalty=1.0,
-        eos_token_id=[151329, 151336, 151338],
-    )
-    gen_kwargs = {**input_ids, **generate_kwargs}
-
-    with torch.no_grad():
-        thread = Thread(target=model.generate, kwargs=gen_kwargs)
-        thread.start()
-        buffer = ""
-        for new_text in streamer:
-            buffer += new_text
-    
-    return buffer
-
-
 MODEL_LIST = ["nikravan/glm-4vq"]
 
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
@@ -252,6 +205,74 @@ EXAMPLES = [
     [{"text": "Quiero armar un JSON, solo el JSON sin texto, que contenga los datos de la primera mitad de la tabla de la imagen (las primeras 10 jurisdicciones 901-910). Ten en cuenta que los valores numéricos son decimales de cuatro dígitos. La tabla contiene las siguientes columnas: Codigo, Nombre, Fecha Inicio, Fecha Cese, Coeficiente Ingresos, Coeficiente Gastos y Coeficiente Unificado. La tabla puede contener valores vacíos, en ese caso dejarlos como null. Cada fila de la tabla representa una jurisdicción con sus respectivos valores.", }]
 ]
 
+# Definir la estructura del mensaje utilizando Pydantic
+class Message(BaseModel):
+    text: str
+    file: Optional[UploadFile] = None
+
+# Definir la función simple_chat
+def simple_chat(message: Message, 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 = []
+
+    # Procesar el mensaje
+    if message.file:
+        file_contents = message.file.file.read()
+        # Aquí deberías procesar el archivo como corresponda, por ejemplo:
+        # choice, contents = mode_load(file_contents)
+        # Por ahora solo agregaremos un marcador de posición
+        choice = "doc"
+        contents = "Contenido del archivo"
+        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)
+    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
+
+    # Configurar parámetros de generación
+    generate_kwargs = dict(
+        max_length=max_length,
+        streamer=streamer,
+        do_sample=True,
+        top_p=top_p,
+        top_k=top_k,
+        temperature=temperature,
+        repetition_penalty=penalty,
+        eos_token_id=[151329, 151336, 151338],
+    )
+    gen_kwargs = {**input_ids, **generate_kwargs}
+
+    # Generar respuesta de manera asíncrona
+    def generate():
+        with torch.no_grad():
+            thread = Thread(target=model.generate, kwargs=gen_kwargs)
+            thread.start()
+            buffer = ""
+            for new_text in streamer:
+                buffer += new_text
+                yield buffer.encode('utf-8')
+
+    return StreamingResponse(generate(), media_type="text/plain")
+
+# Definir la ruta en FastAPI
+@app.post("/chat")
+async def chat(message: Message):
+    return simple_chat(message)
+
 with gr.Blocks(css=CSS, theme="soft", fill_height=True) as demo:
     gr.HTML(TITLE)
     gr.HTML(DESCRIPTION)