File size: 624 Bytes
5e505fa 93e945a 9ce1dd0 5e505fa 766442a 5e505fa 0376814 a89b7b6 5e505fa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 | from transformers import pipeline
from fastapi import FastAPI
from pydantic import BaseModel
app = FastAPI()
generator = pipeline(
'text-generation',
model='microsoft/phi-2',
dtype='auto'
)
class PromptRequest(BaseModel):
prompt: str
system: str = ""
@app.post("/narrate")
def narrate(request: PromptRequest):
formatted = f"Instruct: {request.system}\nInput: {request.prompt}\nOutput:"
result = generator(formatted, max_new_tokens=60, do_sample=True, temperature=0.3)
generated = result[0]['generated_text']
response = generated[len(formatted):]
return {"text": response.strip()} |