Support very small models on CPU

app.py

@@ -4,6 +4,15 @@ from transformer_ranker import TransformerRanker, prepare_popular_models
 
 st.title("Choose Your Transformer")
 
+model_options = {
+    'bert-tiny': 'prajjwal1/bert-tiny',
+    'bert-small': 'prajjwal1/bert-small',
+    'electra-small': 'google/electra-small-discriminator',
+    'deberta-small': 'microsoft/deberta-v3-small',
+    'distilbert-cased': 'distilbert-base-cased',
+    'distilbert-uncased': 'distilbert-base-uncased',
+}
+
 # 1) Select dataset names (from text classification or token classification subcategory)
 dataset_options = ['trec', 'conll2003']  # Example datasets; you can expand this list
 selected_dataset = st.selectbox("Select Dataset", dataset_options)
@@ -13,9 +22,9 @@ downsample_values = [0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
 downsample_ratio = st.select_slider("Dataset Downsample Ratio", options=downsample_values, value=0.2)
 
 # 3) Select multiple models from HuggingFace model hub
-model_options = ['prajjwal1/bert-tiny', 'google/electra-small-discriminator', 'microsoft/deberta-v3-small',
-                 'bert-base-uncased', 'bert-base-cased', 'distilbert-base-uncased', 'roberta-base']  # List of models
-selected_models = st.multiselect("Select Models", model_options, default=model_options[:1])
+model_names = list(model_options.keys())
+selected_models = st.multiselect("Select Models", model_names, default=['bert-tiny', 'electra-small'])
+selected_models = [model_options[model_name] for model_name in selected_models]
 
 # 4) Select the parameter for layer selection with layermean as the default
 layer_options = ['lastlayer', 'layermean', 'bestlayer']