devrhinetum / main.py
samuel-moreira's picture
Update main.py
0f2551f verified
from ctransformers import AutoModelForCausalLM
from fastapi import FastAPI
from pydantic import BaseModel
llm = AutoModelForCausalLM.from_pretrained(
"zephyr-7b-beta.Q4_K_S.gguf",
model_type='llama3',
max_new_tokens=2000,
threads=10,
context_length = 4096)
#Pydantic object
class validation(BaseModel):
prompt: str
#Fast API
app = FastAPI()
@app.post("/llm_on_cpu")
async def stream(item: validation):
print("Start")
system_prompt = """Você lê textos e padroniza em saída JSON.
Leia o currículo e classifique as entidades apropriadamente.
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 do dicionário JSON. 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.]
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)