Skip to content

Commit 972a2ee

Browse files
add in TEXT_TO_TEXT_GENERATION models (#310)
* add in TEXT_TO_TEXT_GENERATION models * checking for dups isn't necessary
1 parent f568dd1 commit 972a2ee

File tree

1 file changed

+4
-2
lines changed
  • llm-service/app/services/caii

1 file changed

+4
-2
lines changed

llm-service/app/services/caii/caii.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -302,9 +302,11 @@ def get_embedding_model(model_name: str) -> BaseEmbedding:
302302

303303

304304
def get_caii_llm_models() -> List[ModelResponse]:
305-
potential_models = get_models_with_task("TEXT_GENERATION")
305+
potential_text_models = get_models_with_task("TEXT_GENERATION")
306+
potential_text_to_text_models = get_models_with_task("TEXT_TO_TEXT_GENERATION")
307+
306308
results: list[Endpoint] = []
307-
for potential in potential_models:
309+
for potential in [*potential_text_models, *potential_text_to_text_models]:
308310
try:
309311
model = get_llm(
310312
endpoint=potential,

0 commit comments

Comments
 (0)