from typing import Annotated, Union
from fastapi import Depends, FastAPI
from fastapi.security import OAuth2PasswordBearer
from pydantic import BaseModel
from transformers import GPT2Tokenizer, TFGPT2LMHeadModel 
app = FastAPI()

oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")


class Item(BaseModel):
    name: str
    description: Union[str,None] = None
    price: float
    tax: Union[float, None] = None

class Prompt(BaseModel):
    text: str


@app.get("/")
async def root():
    return {"message": "Hello world!"}


@app.post("/prompt")
async def prompt_gpt2(prompt:  Prompt):
    tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
    model = TFGPT2LMHeadModel.from_pretrained('gpt2')
    input = tokenizer(prompt.text, return_tensors="tf")
    output = model.generate(**input, max_new_tokens=15, return_dict_in_generate=True, output_scores=True)
    generated_tokens_ids = output.sequences[0]
    return tokenizer.decode(generated_tokens_ids)

    

@app.post("/items/{query}")
async def create_item(query: str, item : Item, token: Annotated[str, Depends(oauth2_scheme)]):
    return item