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