| | 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}") |
| |
|
