mirror of
https://github.com/lllin000/PaperForge.git
synced 2026-07-22 17:00:23 +00:00
feat: urllib-based model download bypasses HF hub + HF_TOKEN field
This commit is contained in:
parent
beda4119e8
commit
f718574087
2 changed files with 80 additions and 12 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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']
|
||||
});
|
||||
|
|
|
|||
Loading…
Reference in a new issue