add data

app.py

@@ -1,6 +1,66 @@
 import gradio as gr
+import huggingface_hub as hfh
+from requests.exceptions import HTTPError
 
+# =====================================================================================================================
+# DATA
+# =====================================================================================================================
+# Dict with the tasks considered in this spaces, {pretty name: space tag}
+TASK_TYPES = {
+    "✍️ Text Generation": "txtgen",
+    "🤏 Summarization": "summ",
+    "🫂 Translation": "trans",
+    "💬 Conversational / Chatbot": "chat",
+    "🤷 Text Question Answering": "txtqa",
+    "🕵️ (Table/Document/Visual) Question Answering": "otherqa",
+    "🎤 Automatic Speech Recognition": "asr",
+    "🌇 Image to Text": "img2txt",
+}
 
+# Dict matching all task types with their possible hub tags, {space tag: (possible hub tags)}
+HUB_TAGS = {
+    "txtgen": ("text-generation", "text2text-generation"),
+    "summ": ("summarization", "text-generation", "text2text-generation"),
+    "trans": ("translation", "text-generation", "text2text-generation"),
+    "chat": ("conversational", "text-generation", "text2text-generation"),
+    "txtqa": ("text-generation", "text2text-generation"),
+    "otherqa": ("table-question-answering", "document-question-answering", "visual-question-answering"),
+    "asr": ("automatic-speech-recognition"),
+    "img2txt": ("image-to-text"),
+}
+assert len(TASK_TYPES) == len(TASK_TYPES)
+assert all(tag in HUB_TAGS for tag in TASK_TYPES.values())
+
+# Dict with the problems considered in this spaces, {problem: space tag}
+PROBLEMS = {
+    "I would like a ChatGPT-like model": "chatgpt",
+    "The output is too repetitive": "repetitive",
+    "The output is too short": "short",
+    "Other": "show_all"
+}
+
+# =====================================================================================================================
+
+
+# =====================================================================================================================
+# SUGGESTIONS LOGIC
+# =====================================================================================================================
+def is_valid_task_for_model(model_name, task_type):
+    if model_name == "":
+        return True
+    try:
+        model_tags = hfh.HfApi().model_info(model_name).tags
+    except HTTPError:
+        return True  # Assume everything is okay
+
+    possible_tags = HUB_TAGS[TASK_TYPES[task_type]]
+    return any(tag in model_tags for tag in possible_tags)
+# =====================================================================================================================
+
+
+# =====================================================================================================================
+# GRADIO
+# =====================================================================================================================
 demo = gr.Blocks()
 with demo:
     gr.Markdown(
@@ -20,37 +80,23 @@ with demo:
     with gr.Row():
         with gr.Column():
             task_type = gr.Dropdown(
-                label="Task Type",
-                choices=[
-                    "✍️ Text Generation",
-                    "🤏 Summarization",
-                    "🫂 Translation",
-                    "💬 Conversational / Chatbot",
-                    "🕵️ (Table/Document/Visual/Text) Question Answering",
-                    "🎤 Automatic Speech Recognition",
-                    "🌇 Image to Text",
-                ],
-                value="Text Generation",
+                label="What is the task you're trying to accomplish?",
+                choices=list(TASK_TYPES.keys()),
+                value=TASK_TYPES[0],
                 interactive=True,
             )
             model_name = gr.Textbox(
-                label="Which model are you using?",
-                placeholder="google/flan-t5-xl",
+                label="Which model are you using? (leave blank if you haven't decided)",
+                placeholder="e.g. google/flan-t5-xl",
                 interactive=True,
             )
             problem_type = gr.Dropdown(
                 label="What is your current problem?",
-                choices=[
-                    "I would like a ChatGPT-like model",
-                    "The output is too repetitive",
-                    "The output is too short or too long",
-                    "Other"
-                ],
-                value="Text Generation",
+                choices=list(PROBLEMS.keys()),
                 interactive=True,
             )
-            get_suggestions = gr.Button(label="Get Suggestions!")
-        with gr.Column(min_width=600):
+            get_suggestions = gr.Button(value="Get Suggestions!")
+        with gr.Column(scale=2):
             box = gr.Textbox()
 
     gr.Markdown(
@@ -59,6 +105,7 @@ with demo:
         [new discussion](https://huggingface.co/spaces/joaogante/generate_quality_improvement/discussions) 🙏
         """
     )
+# =====================================================================================================================
 
 if __name__ == "__main__":
     demo.launch()