Spaces:
Sleeping
Sleeping
| import logging | |
| import threading | |
| import uuid | |
| from typing import Callable, List | |
| import chromadb | |
| import numpy as np | |
| from langchain_chroma import Chroma | |
| from langchain_core.documents import Document | |
| from langchain_core.embeddings import Embeddings | |
| from back_end.config import ( | |
| EMBEDDING_BATCH_SIZE, | |
| EMBEDDING_DEVICE, | |
| EMBEDDING_MAX_SEQ_LENGTH, | |
| EMBEDDING_MODEL_NAME, | |
| EMBEDDING_MODEL_REVISION, | |
| EMBEDDING_ONNX_FILE, | |
| ) | |
| logger = logging.getLogger(__name__) | |
| ProgressCallback = Callable[[int, int], None] | |
| class _SFRCodeEmbeddings(Embeddings): | |
| """Salesforce/SFR-Embedding-Code-400M_R, run with ONNX Runtime. | |
| Uses the official ONNX export shipped in the model repo (pinned commit), so no | |
| remote Python code is executed (`trust_remote_code` is not needed), torch is | |
| not required, and CPU inference is faster. Pooling = CLS token, then L2 norm, | |
| exactly as the model's sentence-transformers config specifies. | |
| """ | |
| # instruction prefix specified by the Salesforce AI Research team | |
| QUERY_INSTRUCTION = "Instruct: Given Code or Text, retrieve relevant content. Query: " | |
| def __init__(self): | |
| import onnxruntime as ort | |
| from huggingface_hub import hf_hub_download | |
| from tokenizers import Tokenizer | |
| def fetch(filename: str) -> str: | |
| return hf_hub_download(EMBEDDING_MODEL_NAME, filename, revision=EMBEDDING_MODEL_REVISION) | |
| self.tokenizer = Tokenizer.from_file(fetch("tokenizer.json")) | |
| self.tokenizer.enable_truncation(max_length=EMBEDDING_MAX_SEQ_LENGTH) | |
| # The shipped tokenizer pads everything to 512; pad to the longest in the batch instead | |
| self.tokenizer.enable_padding(pad_id=self.tokenizer.token_to_id("[PAD]"), pad_token="[PAD]") | |
| available = ort.get_available_providers() | |
| if EMBEDDING_DEVICE != "cpu" and "CUDAExecutionProvider" in available: | |
| providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] | |
| else: | |
| providers = ["CPUExecutionProvider"] | |
| options = ort.SessionOptions() | |
| options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL | |
| logger.info("Loading %s (%s) with %s ...", EMBEDDING_MODEL_NAME, EMBEDDING_ONNX_FILE, providers[0]) | |
| self.session = ort.InferenceSession(fetch(EMBEDDING_ONNX_FILE), sess_options=options, providers=providers) | |
| self.input_names = {i.name for i in self.session.get_inputs()} | |
| logger.info("Embedding model loaded.") | |
| def _encode(self, texts: List[str], progress: ProgressCallback | None = None) -> np.ndarray: | |
| """progress(done_chars, total_chars) is called after every batch and may raise | |
| to abort. Progress is measured in characters because batches are length-sorted, | |
| so counting texts would make early (short) batches look unrealistically fast.""" | |
| if not texts: | |
| return np.zeros((0, 0), dtype=np.float32) | |
| # Sort by length so each batch pads as little as possible, then restore order | |
| order = sorted(range(len(texts)), key=lambda i: len(texts[i])) | |
| result: list[np.ndarray | None] = [None] * len(texts) | |
| total_chars = sum(len(t) for t in texts) or 1 | |
| done_chars = 0 | |
| if progress: | |
| progress(0, total_chars) | |
| for start in range(0, len(texts), EMBEDDING_BATCH_SIZE): | |
| idx = order[start:start + EMBEDDING_BATCH_SIZE] | |
| encodings = self.tokenizer.encode_batch([texts[i] for i in idx]) | |
| feeds = { | |
| "input_ids": np.array([e.ids for e in encodings], dtype=np.int64), | |
| "attention_mask": np.array([e.attention_mask for e in encodings], dtype=np.int64), | |
| } | |
| if "token_type_ids" in self.input_names: | |
| feeds["token_type_ids"] = np.array([e.type_ids for e in encodings], dtype=np.int64) | |
| hidden = self.session.run(["last_hidden_state"], feeds)[0] | |
| cls = hidden[:, 0, :].astype(np.float32) | |
| cls /= np.clip(np.linalg.norm(cls, axis=1, keepdims=True), 1e-12, None) | |
| for row, i in enumerate(idx): | |
| result[i] = cls[row] | |
| done_chars += sum(len(texts[i]) for i in idx) | |
| if progress: | |
| progress(done_chars, total_chars) | |
| return np.stack(result) | |
| def embed_documents(self, texts: List[str], progress: ProgressCallback | None = None) -> List[List[float]]: | |
| return self._encode(texts, progress=progress).tolist() | |
| def embed_query(self, text: str) -> List[float]: | |
| # The query MUST have the exact instruction prefix applied before encoding | |
| return self._encode([self.QUERY_INSTRUCTION + text])[0].tolist() | |
| # Load the model once and share it across sessions (it holds no user data). | |
| # Previously it was reloaded on every repo. | |
| _model_lock = threading.Lock() | |
| _embedding_model: _SFRCodeEmbeddings | None = None | |
| def get_embedding_model() -> _SFRCodeEmbeddings: | |
| global _embedding_model | |
| with _model_lock: | |
| if _embedding_model is None: | |
| _embedding_model = _SFRCodeEmbeddings() | |
| return _embedding_model | |
| def build_vector_db(documents: List[Document], progress: ProgressCallback | None = None) -> Chroma: | |
| """Embed the chunks and store them in a fresh, in-memory vector DB owned by one session. | |
| In-memory (EphemeralClient) + a random collection name means no shared | |
| on-disk database: one visitor can never query or overwrite another's data, | |
| and everything disappears when the session is dropped. `progress` may raise | |
| to abort (too slow / visitor left) before anything is stored. | |
| """ | |
| model = get_embedding_model() | |
| texts = [doc.page_content for doc in documents] | |
| logger.info("Embedding %d chunks...", len(texts)) | |
| all_embeddings = model.embed_documents(texts, progress=progress) | |
| client = chromadb.EphemeralClient( | |
| settings=chromadb.Settings(anonymized_telemetry=False, allow_reset=False) | |
| ) | |
| vector_db = Chroma( | |
| client=client, | |
| collection_name=f"repo_{uuid.uuid4().hex}", | |
| embedding_function=model, | |
| collection_metadata={"hnsw:space": "cosine"}, | |
| ) | |
| CHROMA_BATCH_SIZE = 5000 | |
| for i in range(0, len(texts), CHROMA_BATCH_SIZE): | |
| vector_db._collection.add( | |
| documents=texts[i:i + CHROMA_BATCH_SIZE], | |
| metadatas=[doc.metadata for doc in documents[i:i + CHROMA_BATCH_SIZE]], | |
| embeddings=all_embeddings[i:i + CHROMA_BATCH_SIZE], | |
| ids=[str(uuid.uuid4()) for _ in documents[i:i + CHROMA_BATCH_SIZE]], | |
| ) | |
| return vector_db | |
| def drop_vector_db(vector_db: Chroma | None) -> None: | |
| if vector_db is None: | |
| return | |
| try: | |
| vector_db.delete_collection() | |
| except Exception: | |
| pass | |