Merge pull request #66 from InternLM/gemma

add gemma-7b
InternLM · Mar 6, 2024 · 871b013 · 871b013
2 parents 5cc0b9e + c51acf0
commit 871b013
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 18 deletions.
diff --git a/openaoe/frontend/dist/assets/index-321b9bf8.js b/openaoe/frontend/dist/assets/index-321b9bf8.js
diff --git a/openaoe/frontend/dist/index.html b/openaoe/frontend/dist/index.html
@@ -32,7 +32,7 @@
         }
       }
     </style>
-    <script type="module" crossorigin src="./assets/index-da64670e.js"></script>
+    <script type="module" crossorigin src="./assets/index-321b9bf8.js"></script>
     <link rel="modulepreload" crossorigin href="./assets/vendor-bad6aacb.js">
     <link rel="stylesheet" href="./assets/index-d5c224e8.css">
   </head>

diff --git a/openaoe/frontend/src/config/api-config.ts b/openaoe/frontend/src/config/api-config.ts
@@ -41,6 +41,7 @@ const claude = {
     model: 'claude-1',
     max_tokens: 5000
 };
+
 const baidu = {
     messages: [
         {
@@ -93,6 +94,12 @@ const mistral = {
     stream: true
 };
 
+const gemma = {
+    model: 'gemma',
+    messages: [],
+    stream: true
+};
+
 const API_PREFIX = '';
 
 const API = new Map();
@@ -125,5 +132,9 @@ API.set('mistral', {
     url: `${API_PREFIX}/v1/mistral/v1/mistral/chat`,
     ...mistral
 });
+API.set('gemma', {
+    url: `${API_PREFIX}/v1/google/v1/gemma/chat`,
+    ...gemma
+});
 
 export default API;
diff --git a/openaoe/frontend/src/services/fetch.ts b/openaoe/frontend/src/services/fetch.ts
@@ -19,7 +19,7 @@ export const getPayload = (provider: string, model: string, prompt: string, mess
     const payload = { ...API.get(provider) || API.get(DEFAULT_PROVIDER) };
     delete payload.url;
     payload.model = model;
-    if (['openai', 'internlm', 'gpt-4', 'mistral'].includes(provider)) {
+    if (['openai', 'internlm', 'gpt-4', 'mistral', 'gemma'].includes(provider)) {
         payload.prompt = prompt;
         payload.messages = messages;
     }
@@ -37,23 +37,28 @@ export const getPayload = (provider: string, model: string, prompt: string, mess
         });
     }
     if (provider === 'google') {
-        // TODO: Messages must alternate between authors.
-        const defaultMessage = {
-            content: 'Hi! How can I help you today?',
-            author: '1'
-        };
-        const formatMessage = messages.map((item) => {
-            return {
-                content: item.text,
-                author: item.sender_type === 'user' ? '0' : '1',
+        if (['gemma-7b', 'gemma-2b'].includes(model)) {
+            payload.prompt = prompt;
+            payload.messages = messages;
+        } else {
+            // TODO: Messages must alternate between authors.
+            const defaultMessage = {
+                content: 'Hi! How can I help you today?',
+                author: '1'
             };
-        });
-        formatMessage.push({
-            content: prompt,
-            author: '0',
-        });
-        formatMessage.unshift(defaultMessage);
-        payload.prompt.messages = formatMessage;
+            const formatMessage = messages.map((item) => {
+                return {
+                    content: item.text,
+                    author: item.sender_type === 'user' ? '0' : '1',
+                };
+            });
+            formatMessage.push({
+                content: prompt,
+                author: '0',
+            });
+            formatMessage.unshift(defaultMessage);
+            payload.prompt.messages = formatMessage;
+        }
     }
     if (provider === 'baidu' || provider === 'claude') {
         const formatMessage = messages.map((item) => {