Merge branch 'feat/huggingface-embedding-support' into deploy/dev

langgenius · Sep 20, 2023 · 2694d4c · 2694d4c
2 parents 901fb33 + 7296444
commit 2694d4c
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 23 deletions.
diff --git a/web/app/components/header/account-setting/model-page/configs/huggingface_hub.tsx b/web/app/components/header/account-setting/model-page/configs/huggingface_hub.tsx
@@ -68,25 +68,14 @@ const config: ProviderConfig = {
         ]
       }
       if (v?.huggingfacehub_api_type === 'inference_endpoints') {
-        if (v?.model_type === 'embeddings') {
-          filteredKeys = [
-            'huggingfacehub_api_type',
-            'huggingfacehub_api_token',
-            'model_name',
-            'huggingfacehub_endpoint_url',
-            'model_type',
-          ]
-        }
-        else {
-          filteredKeys = [
-            'huggingfacehub_api_type',
-            'huggingfacehub_api_token',
-            'model_name',
-            'huggingfacehub_endpoint_url',
-            'task_type',
-            'model_type',
-          ]
-        }
+        filteredKeys = [
+          'huggingfacehub_api_type',
+          'huggingfacehub_api_token',
+          'model_name',
+          'huggingfacehub_endpoint_url',
+          'task_type',
+          'model_type',
+        ]
       }
       return filteredKeys.reduce((prev: FormValue, next: string) => {
         prev[next] = v?.[next] || ''
@@ -209,6 +198,25 @@ const config: ProviderConfig = {
           },
         ],
       },
+      {
+        hidden: (value?: FormValue) => !(value?.huggingfacehub_api_type === 'inference_endpoints' && value?.model_type === 'embeddings'),
+        type: 'radio',
+        key: 'task_type',
+        required: true,
+        label: {
+          'en': 'Task',
+          'zh-Hans': 'Task',
+        },
+        options: [
+          {
+            key: 'feature-extraction',
+            label: {
+              'en': 'Feature Extraction',
+              'zh-Hans': 'Feature Extraction',
+            },
+          },
+        ],
+      },
     ],
   },
 }

diff --git a/web/app/components/header/account-setting/model-page/model-modal/Form.tsx b/web/app/components/header/account-setting/model-page/model-modal/Form.tsx
@@ -1,7 +1,7 @@
 import { useEffect, useState } from 'react'
 import type { Dispatch, FC, SetStateAction } from 'react'
 import { useContext } from 'use-context-selector'
-import type { Field, FormValue, ProviderConfigModal } from '../declarations'
+import { type Field, type FormValue, type ProviderConfigModal, ProviderEnum } from '../declarations'
 import { useValidate } from '../../key-validator/hooks'
 import { ValidatingTip } from '../../key-validator/ValidateStatus'
 import { validateModelProviderFn } from '../utils'
@@ -85,10 +85,31 @@ const Form: FC<FormProps> = ({
   }
 
   const handleFormChange = (k: string, v: string) => {
-    if (mode === 'edit' && !cleared)
+    if (mode === 'edit' && !cleared) {
       handleClear({ [k]: v })
-    else
-      handleMultiFormChange({ ...value, [k]: v }, k)
+    }
+    else {
+      const extraValue: Record<string, string> = {}
+      if (
+        (
+          (k === 'model_type' && v === 'embeddings' && value.huggingfacehub_api_type === 'inference_endpoints')
+          || (k === 'huggingfacehub_api_type' && v === 'inference_endpoints' && value.model_type === 'embeddings')
+        )
+        && modelModal?.key === ProviderEnum.huggingface_hub
+      )
+        extraValue.task_type = 'feature-extraction'
+
+      if (
+        (
+          (k === 'model_type' && v === 'text-generation' && value.huggingfacehub_api_type === 'inference_endpoints')
+          || (k === 'huggingfacehub_api_type' && v === 'inference_endpoints' && value.model_type === 'text-generation')
+        )
+        && modelModal?.key === ProviderEnum.huggingface_hub
+      )
+        extraValue.task_type = 'text-generation'
+
+      handleMultiFormChange({ ...value, [k]: v, ...extraValue }, k)
+    }
   }
 
   const handleFocus = () => {