from transformers import pipeline from fastapi import FastAPI from pydantic import BaseModel app = FastAPI() generator = pipeline('text-generation', model='gpt2') class GenerateRequest(BaseModel): inputs: str parameters: dict = {} @app.post("/generate") def generate(request: GenerateRequest): result = generator( request.inputs, max_new_tokens=request.parameters.get('max_new_tokens', 50) ) return result