File size: 445 Bytes
ed7c2d2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f987867
ed7c2d2
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
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