From bc4992782c5645818e28a6bd4f43ea17a713d3d8 Mon Sep 17 00:00:00 2001
From: Igor Gitman <igitman@nvidia.com>
Date: Fri, 23 Feb 2024 15:36:44 -0800
Subject: [PATCH] Add mixtral to mapping for vocab in trt server

Signed-off-by: Igor Gitman <igitman@nvidia.com>
---
 nemo_skills/inference/server/serve_trt.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/nemo_skills/inference/server/serve_trt.py b/nemo_skills/inference/server/serve_trt.py
index b9a371e0c6..64feb6c66d 100644
--- a/nemo_skills/inference/server/serve_trt.py
+++ b/nemo_skills/inference/server/serve_trt.py
@@ -184,6 +184,7 @@ def read_model_name(engine_dir: str):
     name_map = {
         'MistralForCausalLM'.lower(): 'mistral',
         'LlamaForCausalLM'.lower(): 'llama',
+        'MixtralForCausalLM'.lower(): 'mixtral',
     }
     return name_map[name]