from pydantic import BaseModel

from fastapi.middleware.cors import CORSMiddleware
from happytransformer import HappyTextToText, TTSettings
from TextGen import app

class Generate(BaseModel):
    text:str

def generate_text(prompt: str):
    if prompt == "":
        return {"detail": "Please provide a prompt."}
    else:


        happy_tt = HappyTextToText("T5", "vennify/t5-base-grammar-correction")
        
        args = TTSettings(num_beams=5, min_length=1,  max_length=100000)
        
        result = happy_tt.generate_text(f"grammar: {prompt}", args=args)
        return result.text

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.get("/", tags=["Home"])
def api_home():
    return {'detail': 'Welcome to FastAPI TextGen Tutorial!'}

@app.post("/api/generate",  tags=["Generate"], response_model=Generate)
def inference(input_prompt: str):
    return generate_text(prompt=input_prompt)