Update app.py

app.py

@@ -5,6 +5,7 @@ import os, re, json, time, sys, csv, uuid, datetime
 from typing import List, Dict, Any, Optional
 from functools import lru_cache
 from xml.etree import ElementTree as ET
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 
 import numpy as np
 import requests
@@ -15,6 +16,7 @@ ASSETS_DIR = os.environ.get("ASSETS_DIR", "assets")
 FAISS_PATH = os.environ.get("FAISS_PATH", f"{ASSETS_DIR}/index.faiss")
 META_PATH  = os.environ.get("META_PATH",  f"{ASSETS_DIR}/index_meta.filtered.jsonl")
 REL_CONFIG_PATH = os.environ.get("REL_CONFIG_PATH", f"{ASSETS_DIR}/relevance_config.json")
+QUANTIZE = os.environ.get("QUANTIZE", "4bit")  # "none" | "8bit" | "4bit"
 
 # Models
 BASE_MODEL   = os.environ.get("BASE_MODEL", "mistralai/Mistral-7B-Instruct-v0.2")
@@ -189,17 +191,35 @@ if HF_READ_TOKEN:
 if ADAPTER_REPO:
     ADAPTER_PATH = snapshot_download(repo_id=ADAPTER_REPO, allow_patterns=["*"])
 
+# --- LLM load (quantized optional) ---
 dlog("LLM", f"Loading base model: {BASE_MODEL}")
 tokenizer_lm = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=False)
-base_model   = AutoModelForCausalLM.from_pretrained(
-    BASE_MODEL, torch_dtype=dtype, device_map="auto"
-)
+
+if QUANTIZE in {"8bit", "4bit"}:
+    bnb_config = BitsAndBytesConfig(
+        load_in_8bit=(QUANTIZE == "8bit"),
+        load_in_4bit=(QUANTIZE == "4bit"),
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_use_double_quant=True,
+        bnb_4bit_compute_dtype=torch.float16,
+    )
+    base_model = AutoModelForCausalLM.from_pretrained(
+        BASE_MODEL,
+        device_map="auto",
+        quantization_config=bnb_config,
+    )
+else:
+    base_model = AutoModelForCausalLM.from_pretrained(
+        BASE_MODEL,
+        torch_dtype=dtype,
+        device_map="auto",
+    )
 
 dlog("LLM", f"Loading LoRA adapter from: {ADAPTER_PATH}")
-model_lm     = PeftModel.from_pretrained(base_model, ADAPTER_PATH)
-model_lm.to(device)
+model_lm = PeftModel.from_pretrained(base_model, ADAPTER_PATH)
 model_lm.eval()
 
+
 GEN_ARGS_GROUNDED = dict(
     max_new_tokens=MAX_NEW_TOKENS_GROUNDED,
     do_sample=False,
@@ -901,5 +921,8 @@ with gr.Blocks(theme="soft") as demo:
         outputs=[fb_status, feedback_grp],
     )
 
-demo.queue(max_size=32).launch()
+demo.queue(
+    concurrency_count=int(os.environ.get("CONCURRENCY", "2")),
+    max_size=64
+).launch()