Upload app.py

app.py

@@ -12,6 +12,19 @@ import re
 from typing import Tuple, Optional, List, Dict
 from dataclasses import dataclass
 
+os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "0")  # 關閉 Rust 下載器，避免 single-flight/Reqwest 錯
+HF_BASE = "/data/.huggingface" if (os.path.isdir("/data") and os.access("/data", os.W_OK)) \
+          else os.path.join(tempfile.gettempdir(), "huggingface")
+os.environ.setdefault("HF_HOME", HF_BASE)
+os.environ.setdefault("TRANSFORMERS_CACHE", os.path.join(HF_BASE, "transformers"))
+os.environ.setdefault("HF_HUB_CACHE",       os.path.join(HF_BASE, "hub"))
+for _d in (os.environ["HF_HOME"], os.environ["TRANSFORMERS_CACHE"], os.environ["HF_HUB_CACHE"]):
+    os.makedirs(_d, exist_ok=True)
+    
+print(f"🧠 [LOG] 從 Hugging Face 載入模型（將快取/權重放在 {os.environ['HF_HOME']}）：{model_repo_id}")
+local_model, local_tokenizer = load_model_from_repo(model_repo_id)
+print("✅ [LOG] 模型已就緒（本地推論）")
+    
 import fitz  # PyMuPDF (pymupdf)
 import gradio as gr
 
@@ -81,6 +94,57 @@ HF_TOKEN = os.getenv("HF_TOKEN")  # 若模型 gated 需設定
 _hf_tok = None
 _hf_model = None
 
+# ======== 載入 LLM（T4 預設 4-bit + offload 到 /data） ========
+def load_model_from_repo(repo_id: str):
+    # 1) 先預下載到 /data
+    local_dir = _prefetch_repo(repo_id)
+
+    # 2) tokenizer
+    tok = AutoTokenizer.from_pretrained(local_dir, use_fast=True)
+    if tok.pad_token is None:
+        tok.pad_token = tok.eos_token
+
+    # 3) 4-bit 設定（T4 無 bf16，使用 fp16 作為計算 dtype）
+    qconf = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_use_double_quant=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.float16,
+    )
+
+    # 4) offload 資料夾統一到 /data
+    offload_folder = os.path.join(os.environ["HF_HOME"], "offload")
+    os.makedirs(offload_folder, exist_ok=True)
+
+    # 5) 設定記憶體預算
+    max_memory = {}
+    if _torch.cuda.is_available():
+        num_gpus = _torch.cuda.device_count()
+        # 預留一點空間給 KV cache 與其他進程；你可把 14 調成 12 以更保守
+        per_gpu_budget_gib = 14
+        for i in range(num_gpus):
+            max_memory[i] = f"{per_gpu_budget_gib}GiB"
+        # CPU 也給個上限，用於 offload
+        max_memory["cpu"] = f"{max(4, int(psutil.virtual_memory().available // (1024**3) * 0.7))}GiB"
+        device_map = "auto"  # 讓 transformers 自動把部分層分配到各 GPU/CPU
+        compute_dtype = _torch.float16  # T4 沒 bf16，用 fp16 計算
+    else:
+        # 沒有 GPU：全部走 CPU
+        max_memory["cpu"] = f"{max(4, int(psutil.virtual_memory().available // (1024**3) * 0.9))}GiB"
+        device_map = {"": "cpu"}
+        compute_dtype = _torch.float32
+
+    mdl = AutoModelForCausalLM.from_pretrained(
+        local_dir,
+        quantization_config=qconf,
+        device_map="auto",
+        max_memory=max_memory,
+        offload_folder=offload_folder,
+        low_cpu_mem_usage=True,
+        use_safetensors=True,
+    )
+    return mdl, tok
+
 def _ensure_local_model(logs: Optional[List[str]] = None):
     global _hf_tok, _hf_model
     if _hf_tok is not None and _hf_model is not None: