Merge pull request #15 from codefuse-ai/modelcache_localDB_dev

add fasttext for embedding

Author: peng3307165

modelcache/embedding/__init__.py

@@ -3,6 +3,7 @@
 huggingface = LazyImport("huggingface", globals(), "modelcache.embedding.huggingface")
 data2vec = LazyImport("data2vec", globals(), "modelcache.embedding.data2vec")
 llmEmb = LazyImport("llmEmb", globals(), "modelcache.embedding.llmEmb")
+fasttext = LazyImport("fasttext", globals(), "gptcache.embedding.fasttext")
 
 
 def Huggingface(model="sentence-transformers/all-mpnet-base-v2"):
@@ -15,3 +16,7 @@ def Data2VecAudio(model="facebook/data2vec-audio-base-960h"):
 
 def LlmEmb2vecAudio():
     return llmEmb.LlmEmb2Vec()
+
+
+def FastText(model="en", dim=None):
+    return fasttext.FastText(model, dim)

modelcache/embedding/fasttext.py

@@ -0,0 +1,38 @@
+# -*- coding: utf-8 -*-
+"""
+   Alipay.com Inc.
+   Copyright (c) 2004-2023 All Rights Reserved.
+   ------------------------------------------------------
+   File Name : fasttext.py
+   Author : fuhui.phe
+   Create Time : 2023/12/3 15:40
+   Description : description what the main function of this file
+   Change Activity: 
+        version0 : 2023/12/3 15:40 by fuhui.phe  init
+"""
+import numpy as np
+import os
+from modelcache.utils import import_fasttext
+from modelcache.embedding.base import BaseEmbedding
+import_fasttext()
+import fasttext.util
+
+
+class FastText(BaseEmbedding):
+    def __init__(self, model: str = "en", dim: int = None):
+        self.model_path = os.path.abspath(fasttext.util.download_model(model))
+        self.ft = fasttext.load_model(self.model_path)
+
+        if dim:
+            fasttext.util.reduce_model(self.ft, dim)
+        self.__dimension = self.ft.get_dimension()
+
+    def to_embeddings(self, data, **_):
+        assert isinstance(data, str), "Only allow string as input."
+        emb = self.ft.get_sentence_vector(data)
+        return np.array(emb).astype("float32")
+
+    @property
+    def dimension(self):
+        return self.__dimension
+

modelcache/utils/__init__.py

@@ -48,3 +48,7 @@ def import_faiss():
 
 def import_torch():
     _check_library("torch")
+
+
+def import_fasttext():
+    _check_library("fasttext")