initial commit

venv/Lib/site-packages/langchain_community/embeddings/model2vec.py

@@ -0,0 +1,66 @@
+"""Wrapper around model2vec embedding models."""
+
+from typing import List
+
+from langchain_core.embeddings import Embeddings
+
+
+class Model2vecEmbeddings(Embeddings):
+    """Model2Vec embedding models.
+
+    Install model2vec first, run 'pip install -U model2vec'.
+    The github repository for model2vec is : https://github.com/MinishLab/model2vec
+
+    Example:
+        .. code-block:: python
+
+            from langchain_community.embeddings import Model2vecEmbeddings
+
+            embedding = Model2vecEmbeddings("minishlab/potion-base-8M")
+            embedding.embed_documents([
+                "It's dangerous to go alone!",
+                "It's a secret to everybody.",
+            ])
+            embedding.embed_query(
+                "Take this with you."
+            )
+    """
+
+    def __init__(self, model: str):
+        """Initialize embeddings.
+
+        Args:
+            model: Model name.
+        """
+        try:
+            from model2vec import StaticModel
+        except ImportError as e:
+            raise ImportError(
+                "Unable to import model2vec, please install with "
+                "`pip install -U model2vec`."
+            ) from e
+        self._model = StaticModel.from_pretrained(model)
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """Embed documents using the model2vec embeddings model.
+
+        Args:
+            texts: The list of texts to embed.
+
+        Returns:
+            List of embeddings, one for each text.
+        """
+
+        return self._model.encode(texts).tolist()
+
+    def embed_query(self, text: str) -> List[float]:
+        """Embed a query using the model2vec embeddings model.
+
+        Args:
+            text: The text to embed.
+
+        Returns:
+            Embeddings for the text.
+        """
+
+        return self._model.encode(text).tolist()