Spaces:
Sleeping
Sleeping
| import os | |
| import json | |
| import requests | |
| from pathlib import Path | |
| from typing import List, Dict, Any, Optional | |
| from bs4 import BeautifulSoup | |
| # utils から書き込み先と分割関数を取得 | |
| from .utils import ensure_dirs, data_dir, chunk_text | |
| # 依存は遅延ロード(モデル初期化は重いので) | |
| _model = None | |
| def _embedder(): | |
| """ | |
| SentenceTransformer を遅延初期化。 | |
| - すべてのキャッシュは utils.ensure_dirs() 側で /tmp 等の書き込み可パスへ固定済み | |
| - 環境によってはネット/モデルDLが禁止のことがあるため、呼び出し側での強制ウォームアップはしない | |
| """ | |
| global _model | |
| if _model is not None: | |
| return _model | |
| ensure_dirs() | |
| cache_base = data_dir() / "hf_cache" | |
| # ここで環境変数も最終確認(多重防御) | |
| os.environ.setdefault("HF_HOME", str(data_dir() / "hf_home")) | |
| os.environ.setdefault("HUGGINGFACE_HUB_CACHE", str(cache_base)) | |
| os.environ.setdefault("TRANSFORMERS_CACHE", str(cache_base)) | |
| os.environ.setdefault("SENTENCE_TRANSFORMERS_HOME", str(cache_base)) | |
| os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1") | |
| os.environ.setdefault("HF_TOKEN", "") | |
| from sentence_transformers import SentenceTransformer | |
| # ローカル同梱モデルがあれば優先(ネット不可時の対策) | |
| local_model_dir = data_dir() / "models" / "all-MiniLM-L6-v2" | |
| model_name = str(local_model_dir) if local_model_dir.exists() else os.getenv( | |
| "EMBEDDING_MODEL", "sentence-transformers/all-MiniLM-L6-v2" | |
| ) | |
| _model = SentenceTransformer(model_name, cache_folder=str(cache_base)) | |
| return _model | |
| def _write_chunks(rows: List[Dict[str, Any]]) -> int: | |
| """ | |
| chunks.jsonl に追記(既存は破棄せず、簡易に追加) | |
| 1行: {"text": "...", "source": "path_or_url", "meta": {...}} | |
| """ | |
| ensure_dirs() | |
| out = data_dir() / "chunks.jsonl" | |
| with open(out, "a", encoding="utf-8") as f: | |
| for r in rows: | |
| f.write(json.dumps(r, ensure_ascii=False) + "\n") | |
| return len(rows) | |
| def _load_text_from_url(url: str) -> str: | |
| try: | |
| r = requests.get(url, timeout=15) | |
| r.raise_for_status() | |
| html = r.text | |
| soup = BeautifulSoup(html, "html.parser") | |
| # タイトル + 本文テキスト(簡易) | |
| title = (soup.title.string.strip() if soup.title and soup.title.string else "") | |
| text = soup.get_text("\n", strip=True) | |
| return (title + "\n\n" + text).strip() | |
| except Exception: | |
| return "" | |
| def _load_text_from_file(path: Path) -> str: | |
| # テキスト/Markdown想定(PDF等は最小構成では未対応) | |
| try: | |
| with open(path, "r", encoding="utf-8", errors="ignore") as f: | |
| return f.read() | |
| except Exception: | |
| return "" | |
| def index_files_and_urls(file_paths: Optional[List[str]] = None, urls: Optional[List[str]] = None) -> str: | |
| """ | |
| - 受け取ったファイルとURLからテキストを抽出し、チャンク化して chunks.jsonl に追記 | |
| - 依存を最小化(PDF/Officeは最小構成では対象外) | |
| - モデルのウォームアップは実施しない(ネット不可環境で失敗するため) | |
| """ | |
| ensure_dirs() | |
| file_paths = file_paths or [] | |
| urls = urls or [] | |
| added = 0 | |
| rows: List[Dict[str, Any]] = [] | |
| # ファイル | |
| for p in file_paths: | |
| try: | |
| path = Path(p) | |
| txt = _load_text_from_file(path) | |
| for ch in chunk_text(txt): | |
| rows.append({"text": ch, "source": str(path), "meta": {"kind": "file"}}) | |
| except Exception: | |
| continue | |
| # URL | |
| for u in urls: | |
| txt = _load_text_from_url(u) | |
| for ch in chunk_text(txt): | |
| rows.append({"text": ch, "source": u, "meta": {"kind": "url"}}) | |
| if rows: | |
| added = _write_chunks(rows) | |
| # ここでの warmup を削除(または状態だけ返す) | |
| warmed = False # UIに表示するためのダミー値 | |
| return f"indexed_chunks={added}, warmed_up={warmed}" | |