| 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()} |