Create app.py

app.py

@@ -0,0 +1,298 @@
+from fastapi import FastAPI
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    TextIteratorStreamer
+)
+import torch
+import uvicorn
+import threading
+import json
+
+# =========================
+# APP
+# =========================
+
+app = FastAPI()
+
+stop_flags = {}
+
+# =========================
+# MODEL
+# =========================
+
+MODEL_ID = "Qwen/Qwen2.5-Coder-1.5B-Instruct"
+
+print("🚀 Loading Fast Coder Model...")
+
+device = torch.device(
+    "cuda" if torch.cuda.is_available() else "cpu"
+)
+
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+
+# =========================
+# TOKENIZER
+# =========================
+
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_ID,
+    trust_remote_code=True
+)
+
+# =========================
+# MODEL
+# =========================
+
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    trust_remote_code=True,
+    torch_dtype=torch.float16 if device.type == "cuda" else torch.float32
+)
+
+model = model.to(device)
+model.eval()
+
+print(f"✅ Loaded on {device}")
+
+# =========================
+# REQUEST
+# =========================
+
+class ChatRequest(BaseModel):
+    message: str
+    conversation_id: str
+    temperature: float = 0.1
+
+# =========================
+# SYSTEM PROMPT
+# =========================
+
+SYSTEM_PROMPT = """
+You are a strict expert programming assistant.
+
+CRITICAL RULES:
+- Answer ONLY the user's latest request
+- NEVER continue conversations
+- NEVER generate extra examples unless asked
+- NEVER explain unnecessarily
+- NEVER repeat code
+- NEVER simulate dialogue
+- ALWAYS close markdown code blocks properly
+- ALWAYS return complete executable code
+- Stop immediately after final answer
+
+CODE RULES:
+- Use proper markdown
+- Use ```language
+- Keep formatting clean
+- No duplicate code
+- No unfinished code
+"""
+
+# =========================
+# STOP WORDS
+# =========================
+
+STOP_WORDS = [
+    "<|im_end|>",
+    "<|endoftext|>",
+    "<|eot_id|>",
+    "User:",
+    "Assistant:",
+    "Human:"
+]
+
+# =========================
+# CLEAN OUTPUT
+# =========================
+
+def clean_output(text):
+
+    for w in STOP_WORDS:
+
+        if w in text:
+            text = text.split(w)[0]
+
+    return text.strip()
+
+# =========================
+# BUILD INPUTS
+# =========================
+
+def build_inputs(message):
+
+    messages = [
+        {
+            "role": "system",
+            "content": SYSTEM_PROMPT
+        },
+        {
+            "role": "user",
+            "content": message
+        }
+    ]
+
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+
+    return tokenizer(
+        text,
+        return_tensors="pt"
+    ).to(device)
+
+# =========================
+# STOP ENDPOINT
+# =========================
+
+@app.post("/v1/stop")
+def stop(data: dict):
+
+    stop_flags[data.get("conversation_id")] = True
+
+    return {
+        "status": "stopped"
+    }
+
+# =========================
+# NORMAL CHAT
+# =========================
+
+@app.post("/v1/chat")
+def chat(req: ChatRequest):
+
+    inputs = build_inputs(req.message)
+
+    with torch.inference_mode():
+
+        output = model.generate(
+            **inputs,
+            max_new_tokens=512,
+            do_sample=False,
+            temperature=req.temperature,
+            top_p=1.0,
+            repetition_penalty=1.08,
+            pad_token_id=tokenizer.eos_token_id,
+            eos_token_id=tokenizer.eos_token_id
+        )
+
+    result = tokenizer.decode(
+        output[0][inputs.input_ids.shape[1]:],
+        skip_special_tokens=True
+    )
+
+    result = clean_output(result)
+
+    return {
+        "response": result
+    }
+
+# =========================
+# STREAM CHAT
+# =========================
+
+@app.post("/v1/chat/stream")
+def stream_chat(req: ChatRequest):
+
+    inputs = build_inputs(req.message)
+
+    streamer = TextIteratorStreamer(
+        tokenizer,
+        skip_prompt=True,
+        skip_special_tokens=True
+    )
+
+    generation_kwargs = dict(
+        **inputs,
+        streamer=streamer,
+        max_new_tokens=512,
+        do_sample=False,
+        temperature=req.temperature,
+        top_p=1.0,
+        repetition_penalty=1.08,
+        pad_token_id=tokenizer.eos_token_id,
+        eos_token_id=tokenizer.eos_token_id
+    )
+
+    thread = threading.Thread(
+        target=model.generate,
+        kwargs=generation_kwargs
+    )
+
+    thread.start()
+
+    def generate():
+
+        full_text = ""
+
+        for token in streamer:
+
+            if stop_flags.get(req.conversation_id):
+
+                stop_flags[req.conversation_id] = False
+                break
+
+            if not token:
+                continue
+
+            stop_hit = False
+
+            for sw in STOP_WORDS:
+
+                if sw in token:
+                    token = token.split(sw)[0]
+                    stop_hit = True
+                    break
+
+            if token:
+
+                full_text += token
+
+                # stop after completed markdown block
+                if full_text.count("```") >= 2:
+                    yield f"data: {json.dumps({'choices':[{'delta':{'content': token}}]})}\n\n"
+                    break
+
+                yield f"data: {json.dumps({'choices':[{'delta':{'content': token}}]})}\n\n"
+
+            if stop_hit:
+                break
+
+        full_text = clean_output(full_text)
+
+        yield "event: done\ndata: {}\n\n"
+        yield "data: [DONE]\n\n"
+
+    return StreamingResponse(
+        generate(),
+        media_type="text/event-stream"
+    )
+
+# =========================
+# HEALTH
+# =========================
+
+@app.get("/")
+def root():
+
+    return {
+        "status": "Fast Coder Running 🚀"
+    }
+
+# =========================
+# RUN
+# =========================
+
+if __name__ == "__main__":
+
+    uvicorn.run(
+        "app:app",
+        host="0.0.0.0",
+        port=7860
+    )