Update app.py

app.py

@@ -1,7 +1,16 @@
 # app.py — MCP server using DeepSeek via Hugging Face transformers (or fallback)
 # - Put this file next to config.py (see example below)
-# - It loads the model in LOCAL_MODEL (e.g., a DeepSeek HF checkpoint) via transformers.pipeline
-# - If the model cannot be loaded (no transformers / OOM / missing weights), it falls back to a small CPU model or rule-based responder
+# - Supports LOCAL_MODEL values like:
+#     "deepseek/deepseek-r1-0528" or "deepseek/deepseek-r1-0528:free"
+#   If a ":revision" suffix is present, it will be passed as the `revision=` argument
+#   to transformers.from_pretrained so HF validation is satisfied.
+#
+# - It loads the model via transformers.pipeline if available; otherwise falls back
+#   to google/flan-t5-small or rule-based responder.
+#
+# - Developer instruction: when a user pastes a local path (e.g. /mnt/data/...), the
+#   chat handler passes the path unchanged to process_document(); tool invocation
+#   normalization converts file_path -> file_url (file://...) and optionally file_b64.
 
 from mcp.server.fastmcp import FastMCP
 from typing import Optional, List, Tuple, Any, Dict
@@ -39,7 +48,7 @@ try:
         CLIENT_SECRET,
         REFRESH_TOKEN,
         API_BASE,
-        LOCAL_MODEL,      # e.g. "deepseek-ai/deepseek-r1-7b" or smaller/distilled variant
+        LOCAL_MODEL,      # e.g. "deepseek/deepseek-r1-7b" or "deepseek/deepseek-r1-0528:free"
     )
 except Exception as e:
     raise SystemExit(
@@ -125,16 +134,31 @@ def _normalize_local_path_args(args: Any) -> Any:
     return args
 
 # ----------------------------
-# DeepSeek / HF model loader
+# DeepSeek / HF model loader (with revision parsing)
 # ----------------------------
 LLM_PIPELINE = None
 TOKENIZER = None
 LOADED_MODEL_NAME = None
 
+def _parse_model_and_revision(model_string: str) -> Tuple[str, Optional[str]]:
+    """
+    Accepts model strings like:
+      - 'owner/repo'
+      - 'owner/repo:revision'
+    Returns (repo_id, revision_or_none).
+    """
+    if ":" in model_string:
+        repo_id, revision = model_string.split(":", 1)
+        repo_id = repo_id.strip()
+        revision = revision.strip()
+        return repo_id, revision
+    return model_string, None
+
 def init_deepseek_model():
     """
     Try to load LOCAL_MODEL via transformers.pipeline.
-    If loading fails, try a fallback small model (distilgpt2 or flan-t5-small if seq2seq).
+    If a ':revision' is present, pass revision=... to from_pretrained to avoid HF repo-id validation errors.
+    If loading fails, try a fallback small model (flan-t5-small or distilgpt2).
     """
     global LLM_PIPELINE, TOKENIZER, LOADED_MODEL_NAME
 
@@ -149,30 +173,38 @@ def init_deepseek_model():
         return
 
     try:
-        tokenizer_name = LOCAL_TOKENIZER or LOCAL_MODEL
-        model_name = LOCAL_MODEL
-        LOADED_MODEL_NAME = model_name
+        model_string = LOCAL_MODEL
+        repo_id, revision = _parse_model_and_revision(model_string)
+        tokenizer_name = LOCAL_TOKENIZER or repo_id
+        model_name_for_logging = f"{repo_id}" + (f" (rev={revision})" if revision else "")
+        LOADED_MODEL_NAME = model_name_for_logging
 
         # If model looks like seq2seq (T5/flan) use text2text; else causal
         seq2seq_keywords = ["flan", "t5", "seq2seq"]
-        if any(k in model_name.lower() for k in seq2seq_keywords):
-            TOKENIZER = AutoTokenizer.from_pretrained(tokenizer_name, use_fast=True)
-            model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+        if any(k in repo_id.lower() for k in seq2seq_keywords):
+            if revision:
+                TOKENIZER = AutoTokenizer.from_pretrained(tokenizer_name, use_fast=True, revision=revision)
+                model = AutoModelForSeq2SeqLM.from_pretrained(repo_id, revision=revision)
+            else:
+                TOKENIZER = AutoTokenizer.from_pretrained(tokenizer_name, use_fast=True)
+                model = AutoModelForSeq2SeqLM.from_pretrained(repo_id)
             LLM_PIPELINE = pipeline("text2text-generation", model=model, tokenizer=TOKENIZER)
-            logger.info("Loaded seq2seq model: %s", model_name)
+            logger.info("Loaded seq2seq model: %s", model_name_for_logging)
         else:
-            TOKENIZER = AutoTokenizer.from_pretrained(tokenizer_name, use_fast=True)
-            model = AutoModelForCausalLM.from_pretrained(model_name)
+            if revision:
+                TOKENIZER = AutoTokenizer.from_pretrained(tokenizer_name, use_fast=True, revision=revision)
+                model = AutoModelForCausalLM.from_pretrained(repo_id, revision=revision)
+            else:
+                TOKENIZER = AutoTokenizer.from_pretrained(tokenizer_name, use_fast=True)
+                model = AutoModelForCausalLM.from_pretrained(repo_id)
             LLM_PIPELINE = pipeline("text-generation", model=model, tokenizer=TOKENIZER)
-            logger.info("Loaded causal model: %s", model_name)
+            logger.info("Loaded causal model: %s", model_name_for_logging)
 
     except Exception as e:
         logger.error("Failed to load requested model '%s': %s", LOCAL_MODEL, e)
         traceback.print_exc()
         # Try a small CPU-friendly fallback
-        fallback = None
         try:
-            # prefer an instruction-friendly small model if possible
             fallback = "google/flan-t5-small"
             if "flan" in fallback:
                 TOKENIZER = AutoTokenizer.from_pretrained(fallback, use_fast=True)