feat: urllib-based model download bypasses HF hub + HF_TOKEN field

This commit is contained in:
Research Assistant 2026-05-13 21:33:09 +08:00
parent beda4119e8
commit f718574087
2 changed files with 80 additions and 12 deletions

View file

@ -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)

View file

@ -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']
});