API / app.py
misukisu's picture
Update app.py
b78f939 verified
Raw
History Blame Contribute Delete
905 Bytes
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline
app = FastAPI()
pipe = pipeline(
"text-generation",
model="microsoft/phi-3-mini-4k-instruct"
)
class Req(BaseModel):
message: str
def build_prompt(user_msg: str) -> str:
return (
"<|system|>\n"
"You are a helpful, clear, and concise assistant. "
"Give accurate and well-structured answers."
"<|end|>\n"
"<|user|>\n"
f"{user_msg}"
"<|end|>\n"
"<|assistant|>\n"
)
@app.post("/chat")
def chat(req: Req):
prompt = build_prompt(req.message)
out = pipe(
prompt,
max_new_tokens=200,
temperature=0.6,
top_p=0.9,
do_sample=True,
repetition_penalty=1.1
)
text = out[0]["generated_text"]
reply = text.split("<|assistant|>")[-1].strip()
return {"reply": reply}