| import logging |
| import os |
| from dataclasses import dataclass |
| from enum import Enum |
| from typing import Any, Dict, List, Optional |
|
|
| import mteb |
| from sqlitedict import SqliteDict |
|
|
| from pylate import indexes, models, retrieve |
|
|
| |
| logging.basicConfig( |
| level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" |
| ) |
|
|
|
|
| class IndexType(Enum): |
| """Supported index types.""" |
|
|
| PREBUILT = "prebuilt" |
| LOCAL = "local" |
|
|
|
|
| @dataclass |
| class IndexConfig: |
| """Configuration for a search index.""" |
|
|
| name: str |
| type: IndexType |
| path: str |
| description: Optional[str] = None |
|
|
|
|
| class MCPyLate: |
| """Main server class that manages PyLate indexes and search operations.""" |
|
|
| def __init__(self, override: bool = False): |
| self.logger = logging.getLogger(__name__) |
| dataset_name = "leetcode" |
|
|
| model_name = "lightonai/Reason-ModernColBERT" |
| override = override or not os.path.exists( |
| f"indexes/{dataset_name}_{model_name.split('/')[-1]}" |
| ) |
|
|
| self.model = models.ColBERT( |
| model_name_or_path=model_name, |
| ) |
| self.index = indexes.PLAID( |
| override=override, |
| index_name=f"{dataset_name}_{model_name.split('/')[-1]}", |
| ) |
| self.id_to_doc = SqliteDict( |
| f"./indexes/{dataset_name}_{model_name.split('/')[-1]}/id_to_doc.sqlite", |
| outer_stack=False, |
| ) |
|
|
| self.retriever = retrieve.ColBERT(index=self.index) |
| if override: |
| tasks = mteb.get_tasks(tasks=["BrightRetrieval"]) |
| tasks[0].load_data() |
| for doc, doc_id in zip( |
| list(tasks[0].corpus[dataset_name]["standard"].values()), |
| list(tasks[0].corpus[dataset_name]["standard"].keys()), |
| ): |
| self.id_to_doc[doc_id] = doc |
| self.id_to_doc.commit() |
| documents_embeddings = self.model.encode( |
| sentences=list(tasks[0].corpus[dataset_name]["standard"].values()), |
| batch_size=100, |
| is_query=False, |
| show_progress_bar=True, |
| ) |
|
|
| self.index.add_documents( |
| documents_ids=list(tasks[0].corpus[dataset_name]["standard"].keys()), |
| documents_embeddings=documents_embeddings, |
| ) |
| self.logger.info("Created PyLate MCP Server") |
|
|
| def get_document( |
| self, |
| docid: str, |
| ) -> Optional[Dict[str, Any]]: |
| """Retrieve full document by document ID.""" |
|
|
| return {"docid": docid, "text": self.id_to_doc[docid]} |
|
|
| def search(self, query: str, k: int = 10) -> List[Dict[str, Any]]: |
| """Perform multi-vector search on specified index.""" |
| try: |
| query_embeddings = self.model.encode( |
| sentences=[query], |
| is_query=True, |
| show_progress_bar=True, |
| batch_size=32, |
| ) |
| scores = self.retriever.retrieve(queries_embeddings=query_embeddings, k=20) |
| results = [] |
| for score in scores[0]: |
| results.append( |
| { |
| "docid": score["id"], |
| "score": round(score["score"], 5), |
| "text": self.id_to_doc[score["id"]], |
| |
| |
| |
| } |
| ) |
| return results |
| except Exception as e: |
| self.logger.error(f"Search failed: {e}") |
| raise RuntimeError(f"Search operation failed: {e}") |
|
|
