Update utils.py to add llama3 model to together_ai api

google-deepmind · Nov 20, 2024 · 028cc65 · 028cc65
1 parent c61dc78
commit 028cc65
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/concordia/language_model/utils.py b/concordia/language_model/utils.py
@@ -75,7 +75,10 @@ def language_model_setup(
   elif api_type == 'pytorch_gemma':
     cls = pytorch_gemma_model.PyTorchGemmaLanguageModel
   elif api_type == 'together_ai':
-    cls = together_ai.Gemma2
+    if 'llama' in model_name.lower():
+       cls = together_ai.Llama3
+    else:
+       cls = together_ai.Gemma2
   else:
     raise ValueError(f'Unrecognized api type: {api_type}')