Update src/model_loader.py

src/model_loader.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 import os
 import time
 from functools import lru_cache
-from typing import Dict, List
+from typing import Dict, List, Tuple
 
 import torch
 from transformers import (
@@ -12,23 +12,42 @@ from transformers import (
     AutoTokenizer,
     BitsAndBytesConfig,
 )
-from accelerate import init_empty_weights, load_checkpoint_and_dispatch
 
 # --- Diagnostic print to confirm runtime versions ---
 import transformers
 print("[init]", "torch", torch.__version__, "transformers", transformers.__version__)
 
-
 HF_CACHE = os.environ.get("HF_HOME") or os.environ.get("TRANSFORMERS_CACHE") or "/data/econsult/hf_cache"
 
-# Accept MODEL_ID (preferred) or fallback to MODEL_PRIMARY_ID to avoid env-name drift.
-MODEL_PRIMARY_ID = os.environ.get("MODEL_ID") or os.environ.get("MODEL_PRIMARY_ID", "google/medgemma-27b-text-it")
-MODEL_FALLBACK_ID = os.environ.get("MODEL_FALLBACK_ID", "google/medgemma-4b-text-it")
+# -------------------- Env normalization --------------------
+
+def _resolve_model_ids() -> Tuple[str, str]:
+    """
+    Resolve primary/fallback with precedence:
+    - Primary: Model_ID > MODEL_ID > MODEL_PRIMARY_ID > default
+    - Fallback: Model_Fallback_ID > MODEL_FALLBACK_ID > default
+    """
+    env = os.environ
+    primary = (
+        env.get("Model_ID") or
+        env.get("MODEL_ID") or
+        env.get("MODEL_PRIMARY_ID") or
+        "google/medgemma-27b-text-it"
+    )
+    fallback = (
+        env.get("Model_Fallback_ID") or
+        env.get("MODEL_FALLBACK_ID") or
+        "google/medgemma-4b-it"
+    )
+    return primary.strip(), fallback.strip()
+
+def _force_cpu() -> bool:
+    return str(os.environ.get("FORCE_CPU_LLM", "")).strip().lower() in {"1", "true", "yes"}
 
+# -------------------- Device & model selection --------------------
 
 def _pick_device_and_quant() -> Dict[str, object]:
-    cuda = torch.cuda.is_available()
-    if cuda:
+    if torch.cuda.is_available() and not _force_cpu():
         quant = BitsAndBytesConfig(
             load_in_4bit=True,
             bnb_4bit_use_double_quant=True,
@@ -36,20 +55,19 @@ def _pick_device_and_quant() -> Dict[str, object]:
             bnb_4bit_compute_dtype=torch.bfloat16,
         )
         return {"device_map": "auto", "quantization_config": quant, "torch_dtype": torch.bfloat16}
-    else:
-        return {"device_map": {"": "cpu"}, "torch_dtype": torch.float32}
-
-
-@lru_cache(maxsize=1)
-def _select_ids() -> str:
-    # Prefer explicit env override; else keep default.
-    model_id = (os.environ.get("MODEL_ID") or MODEL_PRIMARY_ID).strip()
-    fb = MODEL_FALLBACK_ID.strip()
-    # Simple sanity hints
-    if not model_id:
-        model_id = fb
-    return model_id
-
+    # CPU path
+    return {"device_map": {"": "cpu"}, "torch_dtype": torch.float32}
+
+def _select_runtime_model_id() -> Tuple[str, bool, str]:
+    """
+    Returns (selected_model_id, is_fallback, device_label)
+    device_label in {"GPU","CPU"}
+    """
+    primary, fallback = _resolve_model_ids()
+    on_gpu = torch.cuda.is_available() and not _force_cpu()
+    if on_gpu:
+        return primary, False, "GPU"
+    return fallback, True, "CPU"
 
 @lru_cache(maxsize=1)
 def _load_tokenizer(model_id: str):
@@ -59,39 +77,47 @@ def _load_tokenizer(model_id: str):
         tok.pad_token = tok.eos_token
     return tok
 
-
 @lru_cache(maxsize=1)
-def _load_model(model_id: str, use_quant: bool = True):
+def _load_model(model_id: str, use_quant: bool):
     device_kwargs = _pick_device_and_quant() if use_quant else {"device_map": {"": "cpu"}, "torch_dtype": torch.float32}
-    print(f"[model_loader] Loading model: {model_id} | quant={use_quant} | device_kwargs={list(device_kwargs.keys())}")
-    try:
-        model = AutoModelForCausalLM.from_pretrained(
-            model_id,
-            low_cpu_mem_usage=True,
-            trust_remote_code=True,
-            cache_dir=HF_CACHE,
-            **device_kwargs,
-        )
-        model.eval()
-        return model
-    except Exception as e:
-        # Fallback to smaller model on CPU
-        fb = MODEL_FALLBACK_ID
-        print(f"[model_loader] Primary load failed: {e}\nFalling back to: {fb}")
-        model = AutoModelForCausalLM.from_pretrained(
-            fb,
-            low_cpu_mem_usage=True,
-            trust_remote_code=True,
-            cache_dir=HF_CACHE,
-            device_map={"": "cpu"},
-            torch_dtype=torch.float32,
-        )
-        model.eval()
-        return model
+    print(f"[model_loader] Loading model: {model_id} | device_kwargs={list(device_kwargs.keys())}")
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        low_cpu_mem_usage=True,
+        trust_remote_code=True,
+        cache_dir=HF_CACHE,
+        **device_kwargs,
+    )
+    model.eval()
+    return model
+
+# -------------------- Public helpers --------------------
+
+def active_model_status() -> Dict[str, str | bool]:
+    primary, fallback = _resolve_model_ids()
+    selected, is_fallback, device = _select_runtime_model_id()
+    forced = _force_cpu()
+    return {
+        "primary_id": primary,
+        "fallback_id": fallback,
+        "selected_id": selected,
+        "device": device,
+        "is_fallback": bool(is_fallback or (device == "CPU")),
+        "forced_cpu": forced,
+    }
 
+def generate_chat(
+    messages: List[Dict[str, str]],
+    *,
+    max_new_tokens: int = 700,
+    temperature: float = 0.2,
+    top_p: float = 0.95,
+) -> str:
+    selected_id, is_fallback, device = _select_runtime_model_id()
+    tok = _load_tokenizer(selected_id)
+    model = _load_model(selected_id, use_quant=(device == "GPU"))
 
-def _build_prompt(messages: List[Dict[str, str]]) -> str:
-    """Very simple chat prompt for IT models."""
+    # Very simple chat prompt for IT models.
     sys_msgs = [m["content"] for m in messages if m.get("role") == "system"]
     turns = []
     for m in messages:
@@ -99,21 +125,8 @@ def _build_prompt(messages: List[Dict[str, str]]) -> str:
             turns.append(f"User: {m['content']}")
         elif m.get("role") == "assistant":
             turns.append(f"Assistant: {m['content']}")
-    return (sys_msgs[0] + "\n\n" if sys_msgs else "") + "\n".join(turns) + "\nAssistant:"
-
-
-def generate_chat(
-    messages: List[Dict[str, str]],
-    *,
-    max_new_tokens: int = 700,
-    temperature: float = 0.2,
-    top_p: float = 0.95,
-) -> str:
-    model_id = _select_ids()
-    tok = _load_tokenizer(model_id)
-    model = _load_model(model_id, use_quant=torch.cuda.is_available())
+    prompt = (sys_msgs[0] + "\n\n" if sys_msgs else "") + "\n".join(turns) + "\nAssistant:"
 
-    prompt = _build_prompt(messages)
     inputs = tok(prompt, return_tensors="pt").to(model.device)
 
     gen_kwargs = dict(
@@ -133,8 +146,7 @@ def generate_chat(
     text = tok.decode(out[0], skip_special_tokens=True)
     generated = text[len(prompt):].strip()
 
-    print(f"[model_loader] Generated {max_new_tokens} tokens in {dt:.2f}s (temp={temperature}, top_p={top_p})")
-    print(f"[model_loader] Tokenizer loaded: {model_id} | cache={HF_CACHE}")
+    print(f"[model_loader] Generated <= {max_new_tokens} tokens in {dt:.2f}s (temp={temperature}, top_p={top_p}) | {selected_id} on {device} | fallback={is_fallback}")
     return generated
 
 
@@ -145,3 +157,4 @@ def generate_chat(
 
 
 
+