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