from fastapi import FastAPI
from pydantic import BaseModel
from happytransformer import HappyTextToText, TTSettings

app = FastAPI()

happy_tt = HappyTextToText("T5", "vennify/t5-base-grammar-correction")

class InputText(BaseModel):
    txt: str

@app.post("/correct_grammar/")
def correct_grammar(input_text: InputText):
    args = TTSettings(num_beams=5, min_length=1, max_length=100000)
    corrected_text = happy_tt.generate_text(f"grammar: {input_text.txt}", args=args)
    return {"corrected_text": corrected_text.text}

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