Spaces:
Sleeping
Sleeping
File size: 2,280 Bytes
58a4232 35ac688 58a4232 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from pydantic import BaseModel
import uvicorn
app = FastAPI()
# Allow requests from any origin (so your HTML page can call this API)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["*"],
allow_headers=["*"],
)
# Load model + tokenizer
model_id = "Qwen/Qwen2.5-0.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)
chat_pipeline = pipeline(
"text-generation",
model=model,
tokenizer=tokenizer,
)
SYSTEM_PROMPT = (
"You are Bella, a witty and friendly female AI assistant. "
"Always respond as Bella, never use any other name. "
"Keep your replies short, casual, and to the point."
)
class ChatRequest(BaseModel):
message: str
@app.post("/chat")
async def chat(request: ChatRequest):
# Priming turn anchors the persona much more strongly for small models
messages = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": "Who are you?"},
{"role": "assistant", "content": "I'm Bella! What can I help you with?"},
{"role": "user", "content": request.message},
]
prompt = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
)
outputs = chat_pipeline(
prompt,
max_new_tokens=256,
do_sample=True,
temperature=0.2,
top_p=0.9,
repetition_penalty=1.2,
pad_token_id=tokenizer.eos_token_id,
)
generated_text = outputs[0]["generated_text"]
# Strip off the prompt, leaving only the new reply
response = generated_text[len(prompt):].strip()
# Safety net: cut off if the model starts hallucinating a new turn
for stop_token in ["<|im_start|>", "<|im_end|>", "User:", "system"]:
if stop_token in response:
response = response.split(stop_token)[0].strip()
if not response:
response = "Hmm, I didn't quite catch that — can you rephrase?"
return {"response": response}
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=7860)
|