"""
Project: Build your own ChatGPT using FastAPI
Author: Kaiss Bouali
Tutorial: 
"""

import uvicorn
from fastapi import FastAPI, Request, Form
from transformers import GPT2Tokenizer, GPT2LMHeadModel
from pydantic import BaseModel

model = GPT2LMHeadModel.from_pretrained("gpt2")
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")


class GenerateRequest(BaseModel):
    prompt: str


app = FastAPI()


@app.post("/generate")
def generate(request: Request, data: GenerateRequest):
    model = GPT2LMHeadModel.from_pretrained("gpt2")
    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
    input_ids = tokenizer.encode(data.prompt, return_tensors="pt")
    generated = model.generate(
        input_ids, max_length=1024, do_sample=True, top_p=0.95, top_k=60
    )
    return tokenizer.decode(generated.tolist()[0], skip_special_tokens=True)


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)