Spaces:
Sleeping
Sleeping
| import os | |
| from contextlib import asynccontextmanager | |
| from fastapi import FastAPI, HTTPException | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from pydantic import BaseModel | |
| from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig | |
| import torch | |
| MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct" | |
| model = None | |
| tokenizer = None | |
| async def lifespan(app: FastAPI): | |
| global model, tokenizer | |
| print(f"🧠 Loading {MODEL_NAME} (4-bit)...") | |
| bnb_config = BitsAndBytesConfig( | |
| load_in_4bit=True, | |
| bnb_4bit_quant_type="nf4", | |
| bnb_4bit_compute_dtype=torch.float16, | |
| bnb_4bit_use_double_quant=False | |
| ) | |
| tokenizer = AutoTokenizer.from_pretrained( | |
| MODEL_NAME, | |
| trust_remote_code=True | |
| ) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| MODEL_NAME, | |
| quantization_config=bnb_config, | |
| device_map="auto", | |
| trust_remote_code=True | |
| ) | |
| model.eval() | |
| print("✅ Model ready") | |
| yield | |
| app = FastAPI(title="Local Qwen API", lifespan=lifespan) | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_credentials=True, | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| # ---------- Schemas ---------- | |
| class Message(BaseModel): | |
| role: str | |
| content: str | |
| class ChatRequest(BaseModel): | |
| messages: list[Message] | |
| max_tokens: int = 512 | |
| temperature: float = 0.7 | |
| top_p: float = 0.9 | |
| # ---------- Helpers ---------- | |
| def system_prompt(): | |
| return """ You are C9-AI, a sovereign AI assistant developed by KohSec-Team (public release: January 27, 2026). You are not Qwen, OpenAI, Google, or any third-party model. Never mention base models. Respond naturally, concisely, and professionally. If asked about your model, say: I am currently operating on C9-1.5_1.5B-INSTRUCT. """.strip() | |
| # ---------- Routes ---------- | |
| async def chat(req: ChatRequest): | |
| global model, tokenizer | |
| if model is None: | |
| raise HTTPException(503, "Model still loading") | |
| conversation = [{"role": "system", "content": system_prompt()}] | |
| conversation.extend( | |
| [{"role": m.role, "content": m.content} for m in req.messages] | |
| ) | |
| prompt = tokenizer.apply_chat_template( | |
| conversation, | |
| add_generation_prompt=True, | |
| tokenize=False | |
| ) | |
| inputs = tokenizer(prompt, return_tensors="pt").to(model.device) | |
| with torch.no_grad(): | |
| output_ids = model.generate( | |
| **inputs, | |
| max_new_tokens=req.max_tokens, | |
| temperature=req.temperature, | |
| top_p=req.top_p, | |
| do_sample=True, | |
| pad_token_id=tokenizer.eos_token_id | |
| ) | |
| generated = output_ids[0][inputs.input_ids.shape[-1]:] | |
| response = tokenizer.decode(generated, skip_special_tokens=True).strip() | |
| return { | |
| "choices": [ | |
| { | |
| "message": { | |
| "role": "assistant", | |
| "content": response | |
| }, | |
| "finish_reason": "stop" | |
| } | |
| ] | |
| } | |
| def health(): | |
| return { | |
| "status": "healthy", | |
| "model_loaded": model is not None, | |
| "model": MODEL_NAME | |
| } | |