Update main.py

main.py

@@ -1,43 +1,32 @@
 from ctransformers import AutoModelForCausalLM
 from fastapi import FastAPI
-from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 
-# 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)
 
-# Objeto Pydantic
-class Validation(BaseModel):
+#Pydantic object
+class validation(BaseModel):
     prompt: str
-
-# Inicializar FastAPI
+#Fast API
 app = FastAPI()
 
-# Função para gerar resposta incrementada
-async def generate_response(prompt: str):
+@app.post("/llm_on_cpu")
+async def stream(item: validation):
+    print("Start")
     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>"
-    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 StreamingResponse(generate_response(item.prompt), media_type="text/plain")
+    prompt = f"{system_prompt}{E_INST}\n{item.prompt}{E_INST}"
+    return llm(prompt)