improved error handling

app.py

@@ -5,6 +5,7 @@ In the demo, the user can write a prompt
 """
 
 import streamlit as st
+from torch.cuda import CudaError
 
 from hanlde_form_submit import on_form_submit
 from on_server_start import main as on_server_start_main
@@ -36,7 +37,7 @@ with st.form("request_form"):
     submitted_prompt: str = st.text_area(
         label="Input for the model, It is highly recommended to write an English prompt.",
         help="Enter the prompt for the model. The model will generate a response based on this prompt.",
-        max_chars=16384,
+        max_chars=2048,
     )
 
     submitted: bool = st.form_submit_button(
@@ -48,6 +49,8 @@ with st.form("request_form"):
     if submitted:
         try:
             output = on_form_submit(selected_model_name, output_length, submitted_prompt)
-            st.write(f"Generated text: {output}")
-        except ValueError as e:
+        except CudaError as e:
+            st.error("Out of memory. Please try a smaller model, shorter prompt, or a smaller output length.")
+        except (ValueError, TypeError, RuntimeError) as e:
             st.error(e)
+        st.write(f"Generated text: {output}")

hanlde_form_submit.py

@@ -56,11 +56,24 @@ def on_form_submit(model_name: str, output_length: int, prompt: str) -> str:
     :param output_length: The size of the groups to use.
     :param prompt: The prompt to use.
     :return: The output of the model.
+    :raises ValueError: If the model name is not supported, the output length is <= 0,
+     the prompt is empty or longer than
+     16384 characters, or the output length is not an integer.
+     TypeError: If the output length is not an integer or the prompt is not a string.
+     RuntimeError: If the model is not found.
     """
     if model_name not in SUPPORTED_MODEL_NAMES:
         raise ValueError(f"The selected model {model_name} is not supported."
                          f"Supported models are all the models in:"
                          f" https://huggingface.co/models?pipeline_tag=text-generation&library=pytorch")
+    if output_length <= 0:
+        raise ValueError(f"The output length {output_length} must be > 0.")
+    if len(prompt) == 0:
+        raise ValueError(f"The prompt must not be empty.")
+    if not isinstance(prompt, str):
+        raise ValueError(f"The prompt must be a string.")
+    if not isinstance(output_length, int):
+        raise ValueError(f"The output length must be an integer.")
     pipeline = create_pipeline(
         model_name=model_name,
         group_size=output_length,