add huggingface model (infiniflow#2624)

### What problem does this PR solve? infiniflow#2469 ### Type of change - [x] New Feature (non-breaking change which adds functionality) --------- Co-authored-by: Kevin Hu <[email protected]>
blooo-io · Sep 27, 2024 · 96f56a3 · 96f56a3
1 parent 1b2f66f
commit 96f56a3
Show file tree

Hide file tree

Showing 9 changed files with 120 additions and 13 deletions.
diff --git a/api/apps/llm_app.py b/api/apps/llm_app.py
@@ -155,6 +155,10 @@ def apikey_json(keys):
     elif factory == "LocalAI":
         llm_name = req["llm_name"]+"___LocalAI"
         api_key = "xxxxxxxxxxxxxxx"
+
+    elif factory == "HuggingFace":
+        llm_name = req["llm_name"]+"___HuggingFace"
+        api_key = "xxxxxxxxxxxxxxx"
 
     elif factory == "OpenAI-API-Compatible":
         llm_name = req["llm_name"]+"___OpenAI-API"

diff --git a/conf/llm_factories.json b/conf/llm_factories.json
@@ -2344,6 +2344,13 @@
             "tags": "LLM",
             "status": "1",
             "llm": []
-        }
+        },
+        {
+      "name": "HuggingFace",
+      "logo": "",
+      "tags": "TEXT EMBEDDING",
+      "status": "1",
+      "llm": []
+    }
     ]
 }
diff --git a/rag/llm/__init__.py b/rag/llm/__init__.py
@@ -18,7 +18,7 @@
 from .cv_model import *
 from .rerank_model import *
 from .sequence2txt_model import *
-from .tts_model import * 
+from .tts_model import *
 
 EmbeddingModel = {
     "Ollama": OllamaEmbed,
@@ -46,7 +46,8 @@
     "SILICONFLOW": SILICONFLOWEmbed,
     "Replicate": ReplicateEmbed,
     "BaiduYiyan": BaiduYiyanEmbed,
-    "Voyage AI": VoyageEmbed
+    "Voyage AI": VoyageEmbed,
+    "HuggingFace":HuggingFaceEmbed,
 }
 
 

diff --git a/rag/llm/chat_model.py b/rag/llm/chat_model.py
@@ -1414,3 +1414,4 @@ def chat_streamly(self, system, history, gen_conf):
                 yield ans + "\n**ERROR**: " + str(e)
 
             yield response._chunks[-1].usage_metadata.total_token_count
+
diff --git a/rag/llm/embedding_model.py b/rag/llm/embedding_model.py
@@ -678,3 +678,40 @@ def encode_queries(self, text):
             texts=text, model=self.model_name, input_type="query"
             )
         return np.array(res.embeddings), res.total_tokens
+
+
+class HuggingFaceEmbed(Base):
+    def __init__(self, key, model_name, base_url=None):
+        if not model_name:
+            raise ValueError("Model name cannot be None")
+        self.key = key
+        self.model_name = model_name
+        self.base_url = base_url or "http://127.0.0.1:8080"
+
+    def encode(self, texts: list, batch_size=32):
+        embeddings = []
+        for text in texts:
+            response = requests.post(
+                f"{self.base_url}/embed",
+                json={"inputs": text},
+                headers={'Content-Type': 'application/json'}
+            )
+            if response.status_code == 200:
+                embedding = response.json()
+                embeddings.append(embedding[0])
+            else:
+                raise Exception(f"Error: {response.status_code} - {response.text}")
+        return np.array(embeddings), sum([num_tokens_from_string(text) for text in texts])
+
+    def encode_queries(self, text):
+        response = requests.post(
+            f"{self.base_url}/embed",
+            json={"inputs": text},
+            headers={'Content-Type': 'application/json'}
+        )
+        if response.status_code == 200:
+            embedding = response.json()
+            return np.array(embedding[0]), num_tokens_from_string(text)
+        else:
+            raise Exception(f"Error: {response.status_code} - {response.text}")
+
diff --git a/web/src/assets/svg/llm/huggingface.svg b/web/src/assets/svg/llm/huggingface.svg
diff --git a/web/src/pages/user-setting/constants.tsx b/web/src/pages/user-setting/constants.tsx
@@ -26,4 +26,5 @@ export const LocalLlmFactories = [
   'TogetherAI',
   'Replicate',
   'OpenRouter',
+  'HuggingFace',
 ];
diff --git a/web/src/pages/user-setting/setting-model/constant.ts b/web/src/pages/user-setting/setting-model/constant.ts
@@ -40,6 +40,7 @@ export const IconMap = {
   Anthropic: 'anthropic',
   'Voyage AI': 'voyage',
   'Google Cloud': 'google-cloud',
+  HuggingFace: 'huggingface',
 };
 
 export const BedrockRegionList = [

diff --git a/web/src/pages/user-setting/setting-model/ollama-modal/index.tsx b/web/src/pages/user-setting/setting-model/ollama-modal/index.tsx
@@ -8,6 +8,20 @@ type FieldType = IAddLlmRequestBody & { vision: boolean };
 
 const { Option } = Select;
 
+const llmFactoryToUrlMap = {
+  Ollama: 'https://huggingface.co/docs/text-embeddings-inference/quick_tour',
+  Xinference: 'https://inference.readthedocs.io/en/latest/user_guide',
+  LocalAI: 'https://localai.io/docs/getting-started/models/',
+  'LM-Studio': 'https://lmstudio.ai/docs/basics',
+  'OpenAI-API-Compatible': 'https://platform.openai.com/docs/models/gpt-4',
+  TogetherAI: 'https://docs.together.ai/docs/deployment-options',
+  Replicate: 'https://replicate.com/docs/topics/deployments',
+  OpenRouter: 'https://openrouter.ai/docs',
+  HuggingFace:
+    'https://huggingface.co/docs/text-embeddings-inference/quick_tour',
+};
+type LlmFactory = keyof typeof llmFactoryToUrlMap;
+
 const OllamaModal = ({
   visible,
   hideModal,
@@ -35,7 +49,9 @@ const OllamaModal = ({
 
     onOk?.(data);
   };
-
+  const url =
+    llmFactoryToUrlMap[llmFactory as LlmFactory] ||
+    'https://huggingface.co/docs/text-embeddings-inference/quick_tour';
   return (
     <Modal
       title={t('addLlmTitle', { name: llmFactory })}
@@ -46,11 +62,7 @@ const OllamaModal = ({
       footer={(originNode: React.ReactNode) => {
         return (
           <Flex justify={'space-between'}>
-            <a
-              href={`https://github.com/infiniflow/ragflow/blob/main/docs/guides/deploy_local_llm.mdx`}
-              target="_blank"
-              rel="noreferrer"
-            >
+            <a href={url} target="_blank" rel="noreferrer">
               {t('ollamaLink', { name: llmFactory })}
             </a>
             <Space>{originNode}</Space>
@@ -72,10 +84,16 @@ const OllamaModal = ({
           rules={[{ required: true, message: t('modelTypeMessage') }]}
         >
           <Select placeholder={t('modelTypeMessage')}>
-            <Option value="chat">chat</Option>
-            <Option value="embedding">embedding</Option>
-            <Option value="rerank">rerank</Option>
-            <Option value="image2text">image2text</Option>
+            {llmFactory === 'HuggingFace' ? (
+              <Option value="embedding">embedding</Option>
+            ) : (
+              <>
+                <Option value="chat">chat</Option>
+                <Option value="embedding">embedding</Option>
+                <Option value="rerank">rerank</Option>
+                <Option value="image2text">image2text</Option>
+              </>
+            )}
           </Select>
         </Form.Item>
         <Form.Item<FieldType>