Update main.py

main.py

@@ -1,32 +1,43 @@
 from ctransformers import AutoModelForCausalLM
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
 from pydantic import BaseModel
+import asyncio
 
-
+# Carregar o modelo
 llm = AutoModelForCausalLM.from_pretrained(
     "zephyr-7b-beta.Q4_K_S.gguf",
     model_type='llama3',
     max_new_tokens=2000,
     threads=10,
-    context_length = 4096)
+    context_length=4096
+)
 
-#Pydantic object
-class validation(BaseModel):
+# Objeto Pydantic
+class Validation(BaseModel):
     prompt: str
-#Fast API
+
+# Inicializar FastAPI
 app = FastAPI()
 
-@app.post("/llm_on_cpu")
-async def stream(item: validation):
-    print("Start")
+# Função para gerar resposta incrementada
+async def generate_response(prompt: str):
     system_prompt = """Leia o currículo e classifique as entidades no meu estilo.
 Respeite o formato de saída JSON.
 Retorne APENAS O DICIONÁRIO JSON e nada mais. 
 O nome dos campos sempre deve estar em português do Brasil.
 Respeite o nome dos campos. São eles: Nome, Email, Telefone, Localização, Experiência(e suas subchaves são Cargo, Empresa, Descrição, Data de Início e Data de Término), Formação(e suas subchaves são Formação, Instituição e Data de Conclusão) e Conhecimentos(suas subchaves são Conhecimento e Instituição).
-Geralmente a descrição da experiência vem descrita logo após o cargo. Procure e inclua esses dados no campo Descrição de cada Experiência.]
+Geralmente a descrição da experiência vem descrita logo após o cargo. Procure e inclua esses dados no campo Descrição de cada Experiência.
 Acrescente todas as informações que encontrar, todas devem ser classificadas nos seus respectivos campos. Não deixe passar nada.
 Caso não encontre alguma informação, retorne o campo com Null."""
     E_INST = "</s>"
-    prompt = f"{system_prompt}{E_INST}\n{item.prompt}{E_INST}"
-    return llm(prompt)
+    full_prompt = f"{system_prompt}{E_INST}\n{prompt}{E_INST}"
+    
+    response = ""
+    for token in llm(full_prompt, stream=True):
+        response += token
+        yield token
+
+@app.post("/llm_on_cpu")
+async def stream(item: Validation):
+    print("Start")
+    return generate_response(item.prompt)