Spaces:
Runtime error
Runtime error
| # app.py — FastAPI backend with safe model offloading, sliding-window rate limiting, and async history persistence | |
| import asyncio | |
| import json | |
| import logging | |
| import os | |
| import threading | |
| import time | |
| from contextlib import asynccontextmanager | |
| from pathlib import Path | |
| from typing import Any, Dict, List, Optional | |
| from fastapi import FastAPI, HTTPException, Request | |
| from fastapi.responses import FileResponse, HTMLResponse, JSONResponse | |
| from fastapi.staticfiles import StaticFiles | |
| from peft import PeftModel | |
| from pydantic import BaseModel, Field | |
| import torch | |
| from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig | |
| # ----- Configuration & Paths ----- | |
| logging.basicConfig( | |
| level=logging.INFO, | |
| format="%(asctime)s [%(levelname)s] %(name)s: %(message)s" | |
| ) | |
| logger = logging.getLogger("colabmind") | |
| BASE_MODEL = "deepseek-ai/deepseek-coder-6.7b-base" | |
| ADAPTER_REPO = "Agasthya0/colabmind-coder-6.7b-ml-qlora" | |
| ROOT = Path(__file__).resolve().parent | |
| STATIC_DIR = ROOT / "static" | |
| STATIC_DIR.mkdir(parents=True, exist_ok=True) | |
| HISTORY_PATH = ROOT / "history.jsonl" | |
| OFFLOAD_DIR = Path("/tmp/llm_offload") | |
| OFFLOAD_DIR.mkdir(parents=True, exist_ok=True) | |
| # Global holder for loaded ML components | |
| ml_resources: Dict[str, Any] = {} | |
| # ----- Rate Limiter (Thread-safe sliding window per IP) ----- | |
| RATE_WINDOW = 60 # Seconds | |
| MAX_REQUESTS = 12 # Requests allowed per window | |
| _rate_lock = threading.Lock() | |
| _client_requests: Dict[str, List[float]] = {} | |
| def check_rate_limit(client_ip: str) -> bool: | |
| now = time.time() | |
| with _rate_lock: | |
| timestamps = _client_requests.setdefault(client_ip, []) | |
| # Drop timestamps outside the active window | |
| while timestamps and timestamps[0] <= now - RATE_WINDOW: | |
| timestamps.pop(0) | |
| if len(timestamps) >= MAX_REQUESTS: | |
| return False | |
| timestamps.append(now) | |
| return True | |
| # ----- History Persistence ----- | |
| _history_lock = threading.Lock() | |
| def append_history(prompt: str, code: str) -> None: | |
| entry = { | |
| "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"), | |
| "prompt": prompt, | |
| "code": code, | |
| } | |
| try: | |
| with _history_lock: | |
| with open(HISTORY_PATH, "a", encoding="utf-8") as f: | |
| f.write(json.dumps(entry, ensure_ascii=False) + "\n") | |
| except Exception as e: | |
| logger.warning("Failed to record query to history: %s", e) | |
| # ----- Model Execution ----- | |
| def run_generation(prompt: str, max_new_tokens: int, temperature: float, top_p: float) -> str: | |
| tokenizer = ml_resources["tokenizer"] | |
| model = ml_resources["model"] | |
| if not prompt or not prompt.strip(): | |
| return "⚠️ Empty prompt" | |
| # Multi-device/CPU/Disk offload handling: target first model layer's device | |
| target_device = model.device | |
| if hasattr(model, "hf_device_map") and "" in model.hf_device_map: | |
| target_device = model.hf_device_map[""] | |
| elif torch.cuda.is_available(): | |
| target_device = "cuda:0" | |
| inputs = tokenizer( | |
| prompt, | |
| return_tensors="pt", | |
| truncation=True, | |
| max_length=2048 | |
| ).to(target_device) | |
| # Hugging Face enforces do_sample=False when temperature == 0.0 | |
| do_sample = temperature > 0.0 | |
| generate_kwargs = { | |
| "max_new_tokens": max_new_tokens, | |
| "eos_token_id": tokenizer.eos_token_id, | |
| "pad_token_id": tokenizer.pad_token_id or tokenizer.eos_token_id, | |
| "do_sample": do_sample, | |
| } | |
| if do_sample: | |
| generate_kwargs["temperature"] = temperature | |
| generate_kwargs["top_p"] = top_p | |
| with torch.no_grad(): | |
| output_ids = model.generate(**inputs, **generate_kwargs) | |
| # Strip prompt tokens from output if present | |
| prompt_len = inputs["input_ids"].shape[-1] | |
| response_tokens = output_ids[0][prompt_len:] | |
| output_text = tokenizer.decode(response_tokens, skip_special_tokens=True) | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| return output_text.strip() | |
| # ----- Lifespan for Clean Startup/Shutdown ----- | |
| async def lifespan(app: FastAPI): | |
| logger.info("Initializing 4-bit BitsAndBytes quantization configuration...") | |
| bnb_config = BitsAndBytesConfig( | |
| load_in_4bit=True, | |
| bnb_4bit_quant_type="nf4", | |
| bnb_4bit_use_double_quant=True, | |
| bnb_4bit_compute_dtype=torch.float16, | |
| llm_int8_enable_fp32_cpu_offload=True, | |
| ) | |
| logger.info("Loading tokenizer from %s...", BASE_MODEL) | |
| tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True) | |
| if tokenizer.pad_token is None: | |
| tokenizer.pad_token = tokenizer.eos_token | |
| logger.info("Loading base model with auto-dispatch and disk offload...") | |
| model_base = AutoModelForCausalLM.from_pretrained( | |
| BASE_MODEL, | |
| quantization_config=bnb_config, | |
| device_map="auto", | |
| offload_folder=str(OFFLOAD_DIR), | |
| trust_remote_code=True, | |
| ) | |
| logger.info("Attaching LoRA adapter from %s...", ADAPTER_REPO) | |
| # Explicitly pass offload_folder to prevent accelerate dispatch crashes on disk-offloaded layers | |
| model = PeftModel.from_pretrained( | |
| model_base, | |
| ADAPTER_REPO, | |
| device_map="auto", | |
| offload_folder=str(OFFLOAD_DIR), | |
| ) | |
| model.eval() | |
| ml_resources["tokenizer"] = tokenizer | |
| ml_resources["model"] = model | |
| logger.info("Model pipeline ready for inference.") | |
| yield | |
| ml_resources.clear() | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| logger.info("Model unloaded from memory.") | |
| # ----- Application & Routes ----- | |
| app = FastAPI(title="ColabMind Coder API", lifespan=lifespan) | |
| app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static") | |
| class PredictRequest(BaseModel): | |
| data: List[Any] = Field(..., min_length=1, description="Expects [prompt, max_tokens, temp, top_p]") | |
| class HistoryItem(BaseModel): | |
| prompt: Optional[str] = "" | |
| code: Optional[str] = "" | |
| def _parse_generation_args(data: List[Any]): | |
| prompt = str(data[0]) if len(data) > 0 and data[0] is not None else "" | |
| max_tokens = int(data[1]) if len(data) > 1 and data[1] else 512 | |
| temperature = float(data[2]) if len(data) > 2 and data[2] is not None else 0.2 | |
| top_p = float(data[3]) if len(data) > 3 and data[3] is not None else 0.9 | |
| max_tokens = max(16, min(2048, max_tokens)) | |
| temperature = max(0.0, min(2.0, temperature)) | |
| top_p = max(0.0, min(1.0, top_p)) | |
| return prompt, max_tokens, temperature, top_p | |
| async def homepage(): | |
| index_file = ROOT / "index.html" | |
| if index_file.exists(): | |
| return FileResponse(str(index_file)) | |
| return HTMLResponse("<h3>Upload index.html to the application root directory.</h3>", status_code=404) | |
| async def run_predict(req: PredictRequest, request: Request): | |
| client_ip = request.client.host if request.client else "unknown" | |
| if not check_rate_limit(client_ip): | |
| raise HTTPException(status_code=429, detail="Too many requests — please slow down.") | |
| try: | |
| prompt, max_tokens, temperature, top_p = _parse_generation_args(req.data) | |
| logger.info("Prediction requested | Length: %d | Client: %s", len(prompt), client_ip) | |
| # Offload heavy synchronous PyTorch call to background thread pool | |
| output = await asyncio.to_thread( | |
| run_generation, prompt, max_tokens, temperature, top_p | |
| ) | |
| # Non-blocking history record | |
| asyncio.create_task(asyncio.to_thread(append_history, prompt, output)) | |
| return JSONResponse({"data": [output]}) | |
| except Exception as e: | |
| logger.exception("Inference failed unexpectedly") | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| async def run_explain(req: PredictRequest, request: Request): | |
| client_ip = request.client.host if request.client else "unknown" | |
| if not check_rate_limit(client_ip): | |
| raise HTTPException(status_code=429, detail="Too many requests — please slow down.") | |
| try: | |
| raw_code, max_tokens, temperature, top_p = _parse_generation_args(req.data) | |
| prompt = ( | |
| "### Instruction:\n" | |
| "Explain the following Python code in detail, line-by-line, and highlight potential bugs or improvements.\n\n" | |
| f"### Code:\n{raw_code}\n\n" | |
| "### Explanation:\n" | |
| ) | |
| output = await asyncio.to_thread( | |
| run_generation, prompt, max_tokens, temperature, top_p | |
| ) | |
| return JSONResponse({"data": [output]}) | |
| except Exception as e: | |
| logger.exception("Code explanation run failed") | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| async def route_history_append(payload: HistoryItem): | |
| try: | |
| await asyncio.to_thread(append_history, payload.prompt or "", payload.code or "") | |
| return JSONResponse({"ok": True}) | |
| except Exception as e: | |
| logger.exception("Manual history append failed") | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| async def route_history_clear(): | |
| try: | |
| with _history_lock: | |
| if HISTORY_PATH.exists(): | |
| HISTORY_PATH.unlink() | |
| return JSONResponse({"ok": True}) | |
| except Exception as e: | |
| logger.exception("History clear operation failed") | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| async def health(): | |
| model = ml_resources.get("model") | |
| if model is None: | |
| return {"status": "starting", "device": "uninitialized"} | |
| device_info = getattr(model, "hf_device_map", str(getattr(model, "device", "unknown"))) | |
| return {"status": "ok", "device_distribution": device_info} |