Update app.py

app.py

@@ -38,30 +38,52 @@ def load_model():
         logger.error(f"Failed to load model: {e}")
         return False
 
-def get_card_info(hub_id: str) -> Tuple[str, str]:
+def get_card_info(hub_id: str, repo_type: str = "auto") -> Tuple[str, str]:
     """Get card information from a Hugging Face hub_id."""
     model_exists = False
     dataset_exists = False
     model_text = None
     dataset_text = None
 
-    # Try getting model card
-    try:
-        info = model_info(hub_id)
-        card = ModelCard.load(hub_id)
-        model_exists = True
-        model_text = card.text
-    except Exception as e:
-        logger.debug(f"No model card found for {hub_id}: {e}")
-
-    # Try getting dataset card
-    try:
-        info = dataset_info(hub_id)
-        card = DatasetCard.load(hub_id)
-        dataset_exists = True
-        dataset_text = card.text
-    except Exception as e:
-        logger.debug(f"No dataset card found for {hub_id}: {e}")
+    # Handle based on repo type
+    if repo_type == "auto":
+        # Try getting model card
+        try:
+            info = model_info(hub_id)
+            card = ModelCard.load(hub_id)
+            model_exists = True
+            model_text = card.text
+        except Exception as e:
+            logger.debug(f"No model card found for {hub_id}: {e}")
+
+        # Try getting dataset card
+        try:
+            info = dataset_info(hub_id)
+            card = DatasetCard.load(hub_id)
+            dataset_exists = True
+            dataset_text = card.text
+        except Exception as e:
+            logger.debug(f"No dataset card found for {hub_id}: {e}")
+    elif repo_type == "model":
+        try:
+            info = model_info(hub_id)
+            card = ModelCard.load(hub_id)
+            model_exists = True
+            model_text = card.text
+        except Exception as e:
+            logger.error(f"Failed to get model card for {hub_id}: {e}")
+            raise ValueError(f"Could not find model with id {hub_id}")
+    elif repo_type == "dataset":
+        try:
+            info = dataset_info(hub_id)
+            card = DatasetCard.load(hub_id)
+            dataset_exists = True
+            dataset_text = card.text
+        except Exception as e:
+            logger.error(f"Failed to get dataset card for {hub_id}: {e}")
+            raise ValueError(f"Could not find dataset with id {hub_id}")
+    else:
+        raise ValueError(f"Invalid repo_type: {repo_type}. Must be 'auto', 'model', or 'dataset'")
 
     # Handle different cases
     if model_exists and dataset_exists:
@@ -115,12 +137,12 @@ def generate_summary(card_text: str, card_type: str) -> str:
     """Cached wrapper for generate_summary with TTL."""
     return _generate_summary_gpu(card_text, card_type)
 
-def summarize(hub_id: str = "") -> str:
+def summarize(hub_id: str = "", repo_type: str = "auto") -> str:
     """Interface function for Gradio. Returns JSON format."""
     try:
         if hub_id:
-            # Fetch and infer card type automatically
-            card_type, card_text = get_card_info(hub_id)
+            # Fetch card information with specified repo_type
+            card_type, card_text = get_card_info(hub_id, repo_type)
             
             if card_type == "both":
                 model_text, dataset_text = card_text
@@ -148,7 +170,15 @@ def summarize(hub_id: str = "") -> str:
 def create_interface():
     interface = gr.Interface(
         fn=summarize,
-        inputs=gr.Textbox(label="Hub ID", placeholder="e.g., huggingface/llama-7b"),
+        inputs=[
+            gr.Textbox(label="Hub ID", placeholder="e.g., huggingface/llama-7b"),
+            gr.Radio(
+                choices=["auto", "model", "dataset"], 
+                value="auto", 
+                label="Repository Type",
+                info="Choose 'auto' to detect automatically, or specify the repository type"
+            )
+        ],
         outputs=gr.JSON(label="Output"),
         title="Hugging Face Hub TLDR Generator",
         description="Generate concise summaries of model and dataset cards from the Hugging Face Hub.",
@@ -160,4 +190,4 @@ if __name__ == "__main__":
         interface = create_interface()
         interface.launch()
     else:
-        print("Failed to load model. Please check the logs for details.")
+        print("Failed to load model. Please check the logs for details.")