Initial commit

app.py

@@ -0,0 +1,43 @@
+from typing import Any, Dict, List, Optional
+
+import gradio as gr
+from core import MCPyLate
+from huggingface_hub import snapshot_download
+
+"""
+MCPyLate Server
+A Model Context Protocol server that provides search functionality using PyLate.
+"""
+
+snapshot_download(
+    repo_id="lightonai/leetcode_reasonmoderncolbert",
+    local_dir="indexes/",
+    repo_type="dataset",
+)
+mcpylate = MCPyLate()
+
+
+def pylate_search_leetcode(
+    query: str, k: int = 10, index_name: Optional[str] = None
+) -> List[Dict[str, Any]]:
+    """
+    Search the PyLate with multi-vector models in the leetcode collection containing code problems solutions and return top-k hits
+    Args:
+        query: Search query string
+        k: Number of results to return (default: 10)
+        index_name: Name of index to search (default: use default index)
+    Returns:
+        List of search results with docid, score, text snippet, and index name
+    """
+    return mcpylate.search(query, k)
+
+
+demo = gr.Interface(
+    fn=pylate_search_leetcode,
+    inputs=["text"],
+    outputs="text",
+    title="LeetCode Search",
+    description="Search in leetcode database index using PyLate",
+)
+
+demo.launch(mcp_server=True, share=True)

core.py

@@ -0,0 +1,115 @@
+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
+
+# Configure logging
+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()  # Don't forget to commit to save changes!
+            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"]],
+                        # "text": self.id_to_doc[score["id"]][:200] + "…"
+                        # if len(self.id_to_doc[score["id"]]) > 200
+                        # else 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}")

requirements.txt

@@ -0,0 +1,3 @@
+gradio
+git+https://github.com/lightonai/pylate.git@MCPyLate
+mteb