from fastapi import FastAPI
from fastapi.responses import HTMLResponse
from transformers import pipeline
from fastapi.middleware.cors import CORSMiddleware
app = FastAPI()
app.add_middleware(CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"])
pipe = pipeline(task="text-generation", model="gpt2-large", framework="pt")
@app.get(path="/", response_class=HTMLResponse)
def get_ui():
"""
Returns the HTML page for the UI.
"""
html_content = """
Text Generator
"""
return html_content
@app.get("/generate")
def generate(text: str):
"""
Using the text-generation pipeline from `transformers`, generate text
from the given input text. The model used is `openai-community/gpt2-large`, which
can be found [here]().
"""
output = pipe(text)
return {"output": output[0]["generated_text"]}