add ChatNVIDIA support for google/gemma-2b, google/codegemma-7b, and …

…google/recurrentgemma-2b
langchain-ai · Apr 10, 2024 · 7682885 · 7682885
1 parent 4ce4f6a
commit 7682885
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_statics.py b/libs/ai-endpoints/langchain_nvidia_ai_endpoints/_statics.py
@@ -64,6 +64,12 @@ class Model(BaseModel):
         # 'ai-reranking-4b': {'model_type': 'chat'},
         # 'ai-sdxl-turbo': {'model_type': 'image_out'},
         # 'ai-stable-diffusion-xl-base': {'model_type': 'iamge_out'},
+        "ai-codegemma-7b": {"model_type": "chat", "model_name": "google/codegemma-7b"},
+        "ai-recurrentgemma-2b": {
+            "model_type": "chat",
+            "model_name": "google/recurrentgemma-2b",
+        },
+        "ai-gemma-2b": {"model_type": "chat", "model_name": "google/gemma-2b"},
     }
 )