from ctransformers import AutoModelForCausalLM
from fastapi import FastAPI
from pydantic import BaseModel


llm = AutoModelForCausalLM.from_pretrained("TowerInstruct-7B-v0.2.Q8_0.gguf",
max_new_tokens = 2084,
threads = 3,
)

#Pydantic object
class validation(BaseModel):
    prompt: str
#Fast API
app = FastAPI()

# <|im_start|>user
# Translate the following text from Portuguese into English.
# Portuguese: Um grupo de investigadores lançou um novo modelo para tarefas relacionadas com tradução.
# English:<|im_end|>
# <|im_start|>assistant

@app.post("/translate")
async def stream(item: validation):
    translation_prompt = 'Below is an instruction that describes a task. Write a response that appropriately completes the request.'
    S_INST = "<|im_start|>"
    E_INST = "<|im_end|>"
    user, assistant = "user", "assistant"
    prompt = f"{S_INST}{user}\n{translation_prompt}\nChinese:{item.prompt}\nEnglish:{E_INST}\n{S_INST}{assistant}\n"
    return llm(prompt)