open-notebook/pages/components/model_selector.py
2024-11-13 14:48:00 -03:00

35 lines
783 B
Python

from typing import Literal
import streamlit as st
from open_notebook.domain.models import Model
def model_selector(
label,
key,
selected_id=None,
help=None,
model_type: Literal[
"language", "embedding", "speech_to_text", "text_to_speech"
] = "language",
) -> Model:
models = Model.get_models_by_type(model_type)
models.sort(key=lambda x: (x.provider, x.name))
try:
index = (
next((i for i, m in enumerate(models) if m.id == selected_id), 0)
if selected_id
else 0
)
except Exception:
index = 0
return st.selectbox(
label,
models,
format_func=lambda x: f"{x.provider} - {x.name}",
help=help,
index=index,
key=key,
)