adaptation to new versions of grouped-sampling

hanlde_form_submit.py

@@ -2,12 +2,10 @@ import os
 from time import time
 
 import streamlit as st
-from grouped_sampling import GroupedSamplingPipeLine
+from grouped_sampling import GroupedSamplingPipeLine, is_supported, UnsupportedModelNameException
 
 from download_repo import download_pytorch_model
 from prompt_engeneering import rewrite_prompt
-from supported_models import is_supported, SUPPORTED_MODEL_NAME_PAGES_FORMAT, BLACKLISTED_MODEL_NAMES, \
-    BLACKLISTED_ORGANIZATIONS
 
 
 def is_downloaded(model_name: str) -> bool:
@@ -94,17 +92,10 @@ def on_form_submit(
      TypeError: If the output length is not an integer or the prompt is not a string.
      RuntimeError: If the model is not found.
     """
-    if not is_supported(model_name, 1, 1):
-        raise ValueError(
-            f"The model: {model_name} is not supported."
-            f"The supported models are the models from {SUPPORTED_MODEL_NAME_PAGES_FORMAT}"
-            f" that satisfy the following conditions:\n"
-            f"1. The model has at least one like and one download.\n"
-            f"2. The model is not one of: {BLACKLISTED_MODEL_NAMES}.\n"
-            f"3. The model was not created any of those organizations: {BLACKLISTED_ORGANIZATIONS}.\n"
-        )
     if len(prompt) == 0:
-        raise ValueError(f"The prompt must not be empty.")
+        raise ValueError("The prompt must not be empty.")
+    if not is_supported(model_name):
+        raise UnsupportedModelNameException(model_name)
     st.write(f"Loading model: {model_name}...")
     loading_start_time = time()
     pipeline = create_pipeline(
@@ -114,7 +105,7 @@ def on_form_submit(
     loading_end_time = time()
     loading_time = loading_end_time - loading_start_time
     st.write(f"Finished loading model: {model_name}  in {loading_time:,.2f} seconds.")
-    st.write(f"Generating text...")
+    st.write("Generating text...")
     generation_start_time = time()
     generated_text = generate_text(
         pipeline=pipeline,

tests.py

@@ -1,12 +1,11 @@
 import os
 
 import pytest as pytest
-from grouped_sampling import GroupedSamplingPipeLine
+from grouped_sampling import GroupedSamplingPipeLine, get_full_models_list, UnsupportedModelNameException
 
 from on_server_start import download_useful_models
 from hanlde_form_submit import create_pipeline, on_form_submit
 from prompt_engeneering import rewrite_prompt
-from supported_models import get_supported_model_names
 
 HUGGING_FACE_CACHE_DIR = "/home/yoni/.cache/huggingface/hub"
 
@@ -21,7 +20,7 @@ def test_prompt_engineering():
 
 
 def test_get_supported_model_names():
-    supported_model_names = get_supported_model_names()
+    supported_model_names = get_full_models_list()
     assert len(supported_model_names) > 0
     assert "gpt2" in supported_model_names
     assert all(isinstance(name, str) for name in supported_model_names)
@@ -44,16 +43,13 @@ def test_on_form_submit():
     with pytest.raises(ValueError):
         on_form_submit(model_name, output_length, empty_prompt, web_search=False)
     unsupported_model_name = "unsupported_model_name"
-    with pytest.raises(ValueError):
+    with pytest.raises(UnsupportedModelNameException):
         on_form_submit(unsupported_model_name, output_length, prompt, web_search=False)
 
 
 @pytest.mark.parametrize(
     "model_name",
-    get_supported_model_names(
-        min_number_of_downloads=1000,
-        min_number_of_likes=100,
-    )
+    get_full_models_list()[:3]
 )
 def test_create_pipeline(model_name: str):
     pipeline: GroupedSamplingPipeLine = create_pipeline(model_name, 5)