fix: upd dialogpt en params

assistant_dists/dream/docker-compose.override.yml

@@ -1177,6 +1177,7 @@ services:
         SERVICE_NAME: dialogpt
         PRETRAINED_MODEL_NAME_OR_PATH: microsoft/DialoGPT-medium
         N_HYPOTHESES_TO_GENERATE: 5
+        CONFIG_NAME: dialogpt_en.json
       context: ./services/dialogpt/
     command: flask run -h 0.0.0.0 -p 8125
     environment:

assistant_dists/dream_mini/docker-compose.override.yml

@@ -160,6 +160,7 @@ services:
         SERVICE_NAME: dialogpt
         PRETRAINED_MODEL_NAME_OR_PATH: microsoft/DialoGPT-medium
         N_HYPOTHESES_TO_GENERATE: 5
+        CONFIG_NAME: dialogpt_en.json
       context: ./services/dialogpt/
     command: flask run -h 0.0.0.0 -p 8125
     environment:

services/dialogpt/Dockerfile

@@ -6,6 +6,8 @@ WORKDIR /src
 
 ARG PRETRAINED_MODEL_NAME_OR_PATH
 ENV PRETRAINED_MODEL_NAME_OR_PATH ${PRETRAINED_MODEL_NAME_OR_PATH}
+ARG CONFIG_NAME
+ENV CONFIG_NAME ${CONFIG_NAME}
 ARG SERVICE_PORT
 ENV SERVICE_PORT ${SERVICE_PORT}
 ARG N_HYPOTHESES_TO_GENERATE

services/dialogpt/dialogpt_en.json

@@ -0,0 +1,10 @@
+{
+  "max_length": 50,
+  "temperature": 0.6,
+  "do_sample": true,
+  "repetition_penalty": 1.3,
+  "no_repeat_ngram_size": 2,
+  "top_k": 50,
+  "top_p": 0.95,
+  "num_return_sequences": 3
+}

services/dialogpt/server.py

@@ -1,6 +1,7 @@
 import logging
-import time
+import json
 import os
+import time
 
 import sentry_sdk
 import torch
@@ -15,11 +16,15 @@
 logger = logging.getLogger(__name__)
 
 PRETRAINED_MODEL_NAME_OR_PATH = os.environ.get("PRETRAINED_MODEL_NAME_OR_PATH")
+N_HYPOTHESES_TO_GENERATE = int(os.environ.get("N_HYPOTHESES_TO_GENERATE", 1))
+CONFIG_NAME = os.environ.get("CONFIG_NAME")
 logging.info(f"PRETRAINED_MODEL_NAME_OR_PATH = {PRETRAINED_MODEL_NAME_OR_PATH}")
 DEFAULT_CONFIDENCE = 0.9
-N_HYPOTHESES_TO_GENERATE = int(os.environ.get("N_HYPOTHESES_TO_GENERATE", 1))
 ZERO_CONFIDENCE = 0.0
 MAX_HISTORY_DEPTH = 3
+with open(CONFIG_NAME, "r") as f:
+    generation_params = json.load(f)
+generation_params["num_return_sequences"] = N_HYPOTHESES_TO_GENERATE
 
 try:
     tokenizer = AutoTokenizer.from_pretrained(PRETRAINED_MODEL_NAME_OR_PATH)
@@ -38,26 +43,21 @@
 logging.getLogger("werkzeug").setLevel("WARNING")
 
 
-def generate_response(context, model, tokenizer):
+def generate_responses(context, model, tokenizer):
     encoded_context = []
     for uttr in context[-MAX_HISTORY_DEPTH:]:
-        encoded_context += [tokenizer.encode(uttr + tokenizer.eos_token, return_tensors="pt")]
+        encoded_context += [tokenizer.encode(uttr + " " + tokenizer.eos_token, return_tensors="pt")]
     bot_input_ids = torch.cat(encoded_context, dim=-1)
 
     with torch.no_grad():
         if torch.cuda.is_available():
             bot_input_ids = bot_input_ids.to("cuda")
-        chat_history_ids = model.generate(
-            bot_input_ids,
-            do_sample=True,
-            max_length=100,
-            temperature=0.6,
-            repetition_penalty=1.3,
-            pad_token_id=tokenizer.eos_token_id,
-        )
+        chat_history_ids = model.generate(bot_input_ids, pad_token_id=tokenizer.eos_token_id, **generation_params)
         if torch.cuda.is_available():
             chat_history_ids = chat_history_ids.cpu()
-    return tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1] :][0], skip_special_tokens=True)
+
+    outputs = [tokenizer.decode(x[len(bot_input_ids[0]) :], skip_special_tokens=True) for x in chat_history_ids]
+    return outputs
 
 
 @app.route("/respond", methods=["POST"])
@@ -71,8 +71,8 @@ def respond():
         for context in contexts:
             curr_responses = []
             curr_confidences = []
-            for i in range(N_HYPOTHESES_TO_GENERATE):
-                response = generate_response(context, model, tokenizer)
+            outputs = generate_responses(context, model, tokenizer)
+            for response in outputs:
                 if len(response) > 3:
                     # drop too short responses
                     curr_responses += [response]