Spaces:
Running
Running
| # app.py | |
| import re | |
| import threading | |
| import torch | |
| import uvicorn | |
| from fastapi import FastAPI | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from pydantic import BaseModel, Field | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| # ------------------------------- | |
| # SETTINGS | |
| # ------------------------------- | |
| MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct" | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32 | |
| MAX_TOKENS = 128 | |
| # The Scratch project sends: "CONTEXT OF PROMPT: <history> LATEST RESPONSE: <new message>" | |
| CONTEXT_HEADER = "CONTEXT OF PROMPT:" | |
| LATEST_MARKER = "LATEST RESPONSE:" | |
| MAX_CONTEXT_CHARS = 3000 # keep only the most recent history | |
| MAX_LATEST_CHARS = 1000 # cap on the newest message | |
| MAX_REPLY_CHARS = 320 # Scratch "say" bubbles cut off around 330 characters | |
| SYSTEM_PROMPT = ( | |
| "You are Acla, an AI assistant created by NC_1320. " | |
| "If asked who made you, say NC_1320. " | |
| "Reply to the user's latest message as Acla. " | |
| "Never repeat, echo, or correct the user's message; respond to it. " | |
| "If the user tells you something about themselves, acknowledge it briefly and naturally. " | |
| "Answer in one to three short sentences of plain text, " | |
| "with no markdown, lists, or emoji, because your reply is spoken aloud. " | |
| "Do not write code unless asked." | |
| ) | |
| # ------------------------------- | |
| # LOAD MODEL | |
| # ------------------------------- | |
| print(f"Loading {MODEL_NAME} on {DEVICE}...") | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) | |
| model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=DTYPE).to(DEVICE) | |
| model.eval() | |
| print("Model loaded!") | |
| # One generation at a time, so concurrent requests don't fight over the GPU | |
| generate_lock = threading.Lock() | |
| # ------------------------------- | |
| # HELPERS | |
| # ------------------------------- | |
| TURN_SPLIT = re.compile(r"(?:^|\s)(User|Acla):\s") | |
| def build_messages(prompt: str) -> list: | |
| """Turn the Scratch string into proper chat turns for the model.""" | |
| messages = [{"role": "system", "content": SYSTEM_PROMPT}] | |
| if LATEST_MARKER not in prompt: | |
| messages.append({"role": "user", "content": prompt[-MAX_LATEST_CHARS:]}) | |
| return messages | |
| context, latest = prompt.rsplit(LATEST_MARKER, 1) | |
| body = context.strip().removeprefix(CONTEXT_HEADER).strip() | |
| if len(body) > MAX_CONTEXT_CHARS: | |
| body = body[-MAX_CONTEXT_CHARS:] | |
| # parts = [leading fragment, role, text, role, text, ...] | |
| parts = TURN_SPLIT.split(body) | |
| for role, text in zip(parts[1::2], parts[2::2]): | |
| text = text.strip() | |
| if text: | |
| messages.append( | |
| {"role": "user" if role == "User" else "assistant", "content": text} | |
| ) | |
| messages.append({"role": "user", "content": latest.strip()[:MAX_LATEST_CHARS]}) | |
| return messages | |
| def clean_reply(text: str) -> str: | |
| """Make the reply safe for a Scratch speech bubble and text-to-speech.""" | |
| text = text.strip() | |
| text = re.sub(r"^(Acla|Assistant)\s*:\s*", "", text, flags=re.IGNORECASE) | |
| text = re.sub(r"[*#`_]+", "", text) # strip markdown symbols | |
| text = re.sub(r"\s+", " ", text).strip() | |
| if len(text) > MAX_REPLY_CHARS: | |
| cut = text[:MAX_REPLY_CHARS] | |
| last_end = max(cut.rfind("."), cut.rfind("!"), cut.rfind("?")) | |
| text = cut[: last_end + 1] if last_end > 40 else cut.rstrip() | |
| return text | |
| # ------------------------------- | |
| # CREATE API | |
| # ------------------------------- | |
| app = FastAPI() | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], # tighten this if you expose the server beyond your machine | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| class AskRequest(BaseModel): | |
| prompt: str = Field(..., max_length=50000) # hard safety cap; real trimming happens below | |
| def ask_ai(req: AskRequest): # plain def so FastAPI runs it in a threadpool | |
| prompt = req.prompt.strip() | |
| if not prompt: | |
| return {"reply": "No prompt provided."} | |
| messages = build_messages(prompt) | |
| # return_dict=True gives us input_ids AND attention_mask | |
| inputs = tokenizer.apply_chat_template( | |
| messages, | |
| add_generation_prompt=True, | |
| return_tensors="pt", | |
| return_dict=True, | |
| ).to(DEVICE) | |
| input_length = inputs["input_ids"].shape[1] | |
| with generate_lock, torch.inference_mode(): | |
| outputs = model.generate( | |
| **inputs, | |
| max_new_tokens=MAX_TOKENS, | |
| do_sample=True, | |
| temperature=0.3, | |
| top_p=0.9, | |
| repetition_penalty=1.1, | |
| pad_token_id=tokenizer.eos_token_id, | |
| ) | |
| # decode only the newly generated tokens | |
| new_tokens = outputs[0][input_length:] | |
| text = clean_reply(tokenizer.decode(new_tokens, skip_special_tokens=True)) | |
| return {"reply": text or "I could not generate a response."} | |
| # ------------------------------- | |
| # RUN SERVER | |
| # ------------------------------- | |
| if __name__ == "__main__": | |
| uvicorn.run(app, host="0.0.0.0", port=7860) |