OOPFP / app.py
Cpope3's picture
Update app.py
93e945a verified
raw
history blame contribute delete
624 Bytes
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()}