From f718574087e807cdeee42fe6d090b35c2ddeb5ea Mon Sep 17 00:00:00 2001 From: Research Assistant Date: Wed, 13 May 2026 21:33:09 +0800 Subject: [PATCH] feat: urllib-based model download bypasses HF hub + HF_TOKEN field --- paperforge/memory/vector_db.py | 75 +++++++++++++++++++++++++++++----- paperforge/plugin/main.js | 17 +++++++- 2 files changed, 80 insertions(+), 12 deletions(-) diff --git a/paperforge/memory/vector_db.py b/paperforge/memory/vector_db.py index 8fba0d7a..fdb8514d 100644 --- a/paperforge/memory/vector_db.py +++ b/paperforge/memory/vector_db.py @@ -72,29 +72,84 @@ def get_embedding_model(vault: Path): mode = settings.get("vector_db_mode", "local") if mode == "api": - return None # API mode — embedding done externally + return None model_name = settings.get("vector_db_model", "BAAI/bge-small-en-v1.5") - # Apply HF mirror endpoint via huggingface_hub API (works alongside env var) - hf_endpoint = settings.get("vector_db_hf_endpoint", "") or os.environ.get("HF_ENDPOINT", "") - if hf_endpoint: - try: - from huggingface_hub import set_endpoint - set_endpoint(hf_endpoint) - except Exception: - pass - if _cached_model is not None and _cached_model_name == model_name: return _cached_model ST = _get_st() logger.info("Loading embedding model: %s", model_name) + + hf_endpoint = settings.get("vector_db_hf_endpoint", "") or os.environ.get("HF_ENDPOINT", "") + + if hf_endpoint: + local_path = _download_model_via_mirror(model_name, hf_endpoint) + if local_path and (local_path / "modules.json").exists(): + logger.info("Loading from local mirror copy: %s", local_path) + _cached_model = ST(str(local_path)) + _cached_model_name = model_name + return _cached_model + _cached_model = ST(model_name) _cached_model_name = model_name return _cached_model +def _download_model_via_mirror(model_name: str, mirror: str) -> Path | None: + """Download model files from a mirror URL to a local cache directory. + Bypasses huggingface_hub entirely by using urllib directly.""" + try: + import urllib.request + except Exception: + return None + + mirror = mirror.rstrip("/") + base_url = f"{mirror}/{model_name}/resolve/main" + local_dir = Path.home() / ".cache" / "paperforge" / "models" / model_name.replace("/", "--") + + files = [ + "config.json", "modules.json", "config_sentence_transformers.json", + "sentence_bert_config.json", "special_tokens_map.json", + "tokenizer.json", "tokenizer_config.json", "vocab.txt", + "model.safetensors", "pytorch_model.bin", + "1_Pooling/config.json", + ] + + local_dir.mkdir(parents=True, exist_ok=True) + + # Build headers from HF_TOKEN + hf_token = os.environ.get("HF_TOKEN", "") + headers = {} + if hf_token: + headers["Authorization"] = f"Bearer {hf_token}" + + for f in files: + dest = local_dir / f + if dest.exists() and dest.stat().st_size > 0: + continue + dest.parent.mkdir(parents=True, exist_ok=True) + url = f"{base_url}/{f}" + try: + req = urllib.request.Request(url, headers=headers) + with urllib.request.urlopen(req, timeout=600) as resp: + with open(dest, "wb") as out: + while True: + chunk = resp.read(8192) + if not chunk: + break + out.write(chunk) + except Exception: + pass + + # Return path only if core files exist + has_weights = (local_dir / "model.safetensors").exists() or (local_dir / "pytorch_model.bin").exists() + has_config = (local_dir / "modules.json").exists() and (local_dir / "config.json").exists() + return local_dir if has_weights and has_config else None + return _cached_model + + def embed_paper(vault: Path, zotero_key: str, chunks: list[dict]) -> int: """Embed chunks for one paper and insert into ChromaDB. Returns count.""" collection = get_collection(vault) diff --git a/paperforge/plugin/main.js b/paperforge/plugin/main.js index 5dcab0da..91e0820f 100644 --- a/paperforge/plugin/main.js +++ b/paperforge/plugin/main.js @@ -572,6 +572,7 @@ const DEFAULT_SETTINGS = { vector_db_api_key: '', vector_db_api_base: '', vector_db_hf_endpoint: 'https://hf-mirror.com', + vector_db_hf_token: '', vector_db_last_model: '', frozen_skills: {}, }; @@ -2895,6 +2896,18 @@ class PaperForgeSettingTab extends PluginSettingTab { const current = this.plugin.settings.vector_db_hf_endpoint || 'https://hf-mirror.com'; const isPreset = ['https://hf-mirror.com', 'https://huggingface.co'].includes(current); if (isPreset) customInput.settingEl.style.display = 'none'; + + new Setting(containerEl) + .setName('HF Token') + .setDesc('HuggingFace access token (optional, helps with rate limits and gated models)') + .addText(text => { + text.setPlaceholder('hf_...') + .setValue(this.plugin.settings.vector_db_hf_token || '') + .onChange(value => { + this.plugin.settings.vector_db_hf_token = value; + this.plugin.saveSettings(); + }); + }); } _renderVectorNoDeps(containerEl) { @@ -2917,7 +2930,7 @@ class PaperForgeSettingTab extends PluginSettingTab { const notice = new Notice('Installing chromadb + sentence-transformers...', 0); try { const { exec } = require('child_process'); - const env = Object.assign({}, process.env, { PYTHONIOENCODING: 'utf-8', PYTHONUTF8: '1', HF_ENDPOINT: this.plugin.settings.vector_db_hf_endpoint || 'https://hf-mirror.com' }); + const env = Object.assign({}, process.env, { PYTHONIOENCODING: 'utf-8', PYTHONUTF8: '1', HF_ENDPOINT: this.plugin.settings.vector_db_hf_endpoint || 'https://hf-mirror.com', HF_TOKEN: this.plugin.settings.vector_db_hf_token || '' }); await new Promise((resolve, reject) => { exec(`"${pyResult.path}" -m pip install chromadb sentence-transformers`, { encoding: 'utf-8', timeout: 300000, env: env, @@ -3117,7 +3130,7 @@ class PaperForgeSettingTab extends PluginSettingTab { terminalEl.setText(''); const { spawn } = require('child_process'); - const env = Object.assign({}, process.env, { PYTHONIOENCODING: 'utf-8', PYTHONUTF8: '1', HF_ENDPOINT: this.plugin.settings.vector_db_hf_endpoint || 'https://hf-mirror.com' }); + const env = Object.assign({}, process.env, { PYTHONIOENCODING: 'utf-8', PYTHONUTF8: '1', HF_ENDPOINT: this.plugin.settings.vector_db_hf_endpoint || 'https://hf-mirror.com', HF_TOKEN: this.plugin.settings.vector_db_hf_token || '' }); const child = spawn(pyResult.path, ['-m', 'paperforge', '--vault', vp, 'embed', 'build', '--force'], { env: env, stdio: ['ignore', 'pipe', 'pipe'] });