Update app.py

app.py

@@ -21,18 +21,29 @@ TASK_ALIASES = {
 def normalize_task(task):
     return TASK_ALIASES.get(task.lower(), task)
 
+# 
 # 🔍 Extract task from user input
 def extract_task(user_input):
     prompt = (
-        "Given a user query, extract the most likely machine learning task "
-        "from the following list: text-classification, token-classification, "
-        "translation, summarization, question-answering, object-detection. "
-        f"Query: {user_input}. Only return the task name."
+        "You are a helpful AI assistant. Your job is to identify the correct Hugging Face model task "
+        "based on the user's request. Choose one from the following official tasks:\n"
+        "text-classification, token-classification, translation, summarization, object-detection, "
+        "image-classification, question-answering, zero-shot-classification, conversational.\n\n"
+        "Here are some examples:\n"
+        "- 'Translate French to English' → translation\n"
+        "- 'Summarize this legal document' → summarization\n"
+        "- 'Detect humans in images' → object-detection\n"
+        "- 'Classify news articles about finance' → text-classification\n"
+        "- 'Extract people and organizations from legal documents' → token-classification\n"
+        "- 'Build a chatbot for tourists' → conversational\n\n"
+        f"User request: '{user_input}'\n"
+        "Return only the task name from the list above."
     )
     result = task_extractor(prompt, max_new_tokens=10)
     task = result[0]["generated_text"].strip().lower()
     return normalize_task(task)
 
+
 # 🔍 Scrape models from Hugging Face
 def get_models_for_task(task):
     url = f"https://huggingface.co/models?pipeline_tag={task}"