from ctransformers import AutoModelForCausalLM
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
)

# Objeto Pydantic
class Validation(BaseModel):
    prompt: str

# Inicializar FastAPI
app = FastAPI()

# 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.
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 generate_response(item.prompt)