BarinkDev/LargeLanguageModels

1from typing import Annotated, Union2from fastapi import Depends, FastAPI3from fastapi.security import OAuth2PasswordBearer4from pydantic import BaseModel5from transformers import GPT2Tokenizer, TFGPT2LMHeadModel 6app = FastAPI()7 8oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")9 10 11class Item(BaseModel):12    name: str13    description: Union[str,None] = None14    price: float15    tax: Union[float, None] = None16 17class Prompt(BaseModel):18    text: str19 20 21@app.get("/")22async def root():23    return {"message": "Hello world!"}24 25 26@app.post("/prompt")27async def prompt_gpt2(prompt:  Prompt):28    tokenizer = GPT2Tokenizer.from_pretrained('gpt2')29    model = TFGPT2LMHeadModel.from_pretrained('gpt2')30    input = tokenizer(prompt.text, return_tensors="tf")31    output = model.generate(**input, max_new_tokens=15, return_dict_in_generate=True, output_scores=True)32    generated_tokens_ids = output.sequences[0]33    return tokenizer.decode(generated_tokens_ids)34 35    36 37@app.post("/items/{query}")38async def create_item(query: str, item : Item, token: Annotated[str, Depends(oauth2_scheme)]):39    return item40 41