Fix schema for cohere rerank

ystoneman · May 9, 2024 · 8fad532 · 8fad532
1 parent 69cdc88
commit 8fad532
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/lib/shared/layers/python-sdk/python/genai_core/cross_encoder.py b/lib/shared/layers/python-sdk/python/genai_core/cross_encoder.py
@@ -31,6 +31,10 @@ def get_cross_encoder_models():
     if not SAGEMAKER_RAG_MODELS_ENDPOINT:
         models = list(filter(lambda x: x["provider"] != "sagemaker", models))
 
+    for model in models:
+        if 'default' not in model:
+            model['default'] = False
+
     return models