From 41ab90bb8edba02ce2027fd9842578f9dbb7d01c Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Mon, 24 Jun 2024 19:16:20 +0300
Subject: [PATCH] Support more parameters for Infermatic

---
 public/index.html                          | 10 +++++-----
 public/script.js                           |  3 ++-
 public/scripts/textgen-settings.js         |  6 ++++--
 src/constants.js                           | 11 +++++++++++
 src/endpoints/backends/text-completions.js |  2 +-
 5 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/public/index.html b/public/index.html
index b5a2eda86..fcf7687d0 100644
--- a/public/index.html
+++ b/public/index.html
@@ -1141,7 +1141,7 @@
                                         <div class="fa-solid fa-circle-info opacity50p" title="Customize displayed samplers or add custom samplers." data-i18n="[title]Customize displayed samplers or add custom samplers."></div>
                                     </small>
                                 </div>
-                                <div data-newbie-hidden data-tg-type="mancer, vllm, aphrodite, tabby" class="flex-container flexFlowColumn alignitemscenter flexBasis100p flexGrow flexShrink gap0">
+                                <div data-newbie-hidden data-tg-type="mancer, vllm, aphrodite, tabby, infermaticai" class="flex-container flexFlowColumn alignitemscenter flexBasis100p flexGrow flexShrink gap0">
                                     <small data-i18n="Multiple swipes per generation">Multiple swipes per generation</small>
                                     <input type="number" id="n_textgenerationwebui" class="text_pole textAlignCenter" min="1" value="1" step="1" />
                                 </div>
@@ -1263,7 +1263,7 @@
                                         <input class="neo-range-slider" type="range" id="skew_textgenerationwebui" name="volume" min="-5" max="5" step="0.01" />
                                         <input class="neo-range-input" type="number" min="-5" max="5" step="0.01" data-for="skew_textgenerationwebui" id="skew_counter_textgenerationwebui">
                                     </div>
-                                    <div data-newbie-hidden data-tg-type="mancer, ooba, tabby, dreamgen" class="alignitemscenter flex-container flexFlowColumn flexBasis30p flexGrow flexShrink gap0">
+                                    <div data-newbie-hidden data-tg-type="mancer, ooba, tabby, dreamgen, infermaticai" class="alignitemscenter flex-container flexFlowColumn flexBasis30p flexGrow flexShrink gap0">
                                         <small data-i18n="Min Length">Min Length</small>
                                         <input class="neo-range-slider" type="range" id="min_length_textgenerationwebui" name="volume" min="0" max="2000" step="1" />
                                         <input class="neo-range-input" type="number" min="0" max="2000" step="1" data-for="min_length_textgenerationwebui" id="min_length_counter_textgenerationwebui">
@@ -1446,7 +1446,7 @@
                                                     <div class="fa-solid fa-circle-info opacity50p " data-i18n="[title]Ban the eos_token. This forces the model to never end the generation prematurely" title="Ban the eos_token. This forces the model to never end the generation prematurely."></div>
                                                 </label>
                                             </label>
-                                            <label data-tg-type="vllm, aphrodite" class="checkbox_label" for="ignore_eos_token_textgenerationwebui">
+                                            <label data-tg-type="vllm, aphrodite, infermaticai" class="checkbox_label" for="ignore_eos_token_textgenerationwebui">
                                                 <input type="checkbox" id="ignore_eos_token_textgenerationwebui" />
                                                 <small data-i18n="Ignore EOS Token">Ignore EOS Token
                                                     <div class="fa-solid fa-circle-info opacity50p " data-i18n="[title]Ignore the EOS Token even if it generates." title="Ignore the EOS Token even if it generates."></div>
@@ -1471,13 +1471,13 @@
                                                 </label>
                                             </label>
 
-                                            <label data-tg-type="vllm, aphrodite" class="checkbox_label" for="spaces_between_special_tokens_textgenerationwebui">
+                                            <label data-tg-type="vllm, aphrodite, infermaticai" class="checkbox_label" for="spaces_between_special_tokens_textgenerationwebui">
                                                 <input type="checkbox" id="spaces_between_special_tokens_textgenerationwebui" />
                                                 <small data-i18n="Spaces Between Special Tokens">Spaces Between Special Tokens</small>
                                             </label>
                                         </div>
                                     </div>
-                                    <div data-tg-type="mancer, ooba, koboldcpp, vllm, aphrodite, llamacpp, ollama" data-newbie-hidden class="flex-container flexFlowColumn alignitemscenter flexBasis48p flexGrow flexShrink gap0">
+                                    <div data-tg-type="mancer, ooba, koboldcpp, vllm, aphrodite, llamacpp, ollama, infermaticai" data-newbie-hidden class="flex-container flexFlowColumn alignitemscenter flexBasis48p flexGrow flexShrink gap0">
                                         <small data-i18n="Seed" class="textAlignCenter">Seed</small>
                                         <input type="number" id="seed_textgenerationwebui" class="text_pole textAlignCenter" min="-1" value="-1" maxlength="100" />
                                     </div>
diff --git a/public/script.js b/public/script.js
index 62c8eb23a..e276d0af2 100644
--- a/public/script.js
+++ b/public/script.js
@@ -5032,6 +5032,7 @@ function parseAndSaveLogprobs(data, continueFrom) {
                     logprobs = data?.completion_probabilities?.map(x => parseTextgenLogprobs(x.content, [x])) || null;
                 } break;
                 case textgen_types.VLLM:
+                case textgen_types.INFERMATICAI:
                 case textgen_types.APHRODITE:
                 case textgen_types.MANCER:
                 case textgen_types.TABBY: {
@@ -5088,7 +5089,7 @@ function extractMultiSwipes(data, type) {
         return swipes;
     }
 
-    if (main_api === 'openai' || (main_api === 'textgenerationwebui' && [MANCER, VLLM, APHRODITE, TABBY].includes(textgen_settings.type))) {
+    if (main_api === 'openai' || (main_api === 'textgenerationwebui' && [MANCER, VLLM, APHRODITE, TABBY, INFERMATICAI].includes(textgen_settings.type))) {
         if (!Array.isArray(data.choices)) {
             return swipes;
         }
diff --git a/public/scripts/textgen-settings.js b/public/scripts/textgen-settings.js
index bb40d2871..6623f1d9e 100644
--- a/public/scripts/textgen-settings.js
+++ b/public/scripts/textgen-settings.js
@@ -905,6 +905,7 @@ export function parseTextgenLogprobs(token, logprobs) {
         case VLLM:
         case APHRODITE:
         case MANCER:
+        case INFERMATICAI:
         case OOBA: {
             /** @type {Record<string, number>[]} */
             const topLogprobs = logprobs.top_logprobs;
@@ -1020,7 +1021,7 @@ export function isJsonSchemaSupported() {
 }
 
 function getLogprobsNumber() {
-    if (settings.type === VLLM) {
+    if (settings.type === VLLM || settings.type === INFERMATICAI) {
         return 5;
     }
 
@@ -1124,7 +1125,7 @@ export function getTextGenGenerationData(finalPrompt, maxTokens, isImpersonate,
         'best_of': canMultiSwipe ? settings.n : 1,
         'ignore_eos': settings.ignore_eos_token,
         'spaces_between_special_tokens': settings.spaces_between_special_tokens,
-        'seed': settings.seed,
+        'seed': settings.seed >= 0 ? settings.seed : undefined,
     };
     const aphroditeParams = {
         'n': canMultiSwipe ? settings.n : 1,
@@ -1162,6 +1163,7 @@ export function getTextGenGenerationData(finalPrompt, maxTokens, isImpersonate,
 
     switch (settings.type) {
         case VLLM:
+        case INFERMATICAI:
             params = Object.assign(params, vllmParams);
             break;
 
diff --git a/src/constants.js b/src/constants.js
index 157f198d1..bb992530b 100644
--- a/src/constants.js
+++ b/src/constants.js
@@ -224,6 +224,17 @@ const INFERMATICAI_KEYS = [
     'repetition_penalty',
     'stream',
     'stop',
+    'presence_penalty',
+    'frequency_penalty',
+    'min_p',
+    'seed',
+    'ignore_eos',
+    'n',
+    'best_of',
+    'min_tokens',
+    'spaces_between_special_tokens',
+    'skip_special_tokens',
+    'logprobs',
 ];
 
 // https://dreamgen.com/docs/api#openai-text
diff --git a/src/endpoints/backends/text-completions.js b/src/endpoints/backends/text-completions.js
index 49324e443..a6c55acbd 100644
--- a/src/endpoints/backends/text-completions.js
+++ b/src/endpoints/backends/text-completions.js
@@ -340,7 +340,7 @@ router.post('/generate', jsonParser, async function (request, response) {
 
                 // Map InfermaticAI response to OAI completions format
                 if (apiType === TEXTGEN_TYPES.INFERMATICAI) {
-                    data['choices'] = (data?.choices || []).map(choice => ({ text: choice?.message?.content || choice.text }));
+                    data['choices'] = (data?.choices || []).map(choice => ({ text: choice?.message?.content || choice.text, logprobs: choice?.logprobs, index: choice?.index }));
                 }
 
                 return response.send(data);