import torch | |
from fastapi import FastAPI | |
from pydantic import BaseModel | |
app = FastAPI() | |
class InferenceRequest(BaseModel): | |
prompt: str | |
max_length: int = 16 | |
temperature: int = 0.9 | |
def predict(request: InferenceRequest): | |
return { | |
'generated_text': request.prompt | |
} |