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):
    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.]
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>"
    user = "<|user|>"
    prompt = f"{system_prompt}{E_INST}\n{user}\n{item.prompt}{E_INST}"
    return llm(prompt)