from fastapi import FastAPI
from transformers import pipeline
# from huggingface_hub import notebook_login
# notebook_login()


app=FastAPI()

import transformers
import torch

model_id = "google/flan-t5-base"

pipeline = transformers.pipeline(
    "text-generation", model=model_id, model_kwargs={"torch_dtype": torch.bfloat16}, device_map="auto"
)

@app.get("/")
def home():
    return ("Hello world")

@app.get("/generate")
def generate(text:str):
    #call pipeline to generate text
    output=pipeline(text)
    #return the output
    return {"output":output[0]['generated_text']}