Update app.py

app.py

@@ -10,6 +10,7 @@ from cachetools import TTLCache
 from cachetools.func import ttl_cache
 import time
 import os
+import json
 os.environ['HF_TRANSFER'] = "1"
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -37,22 +38,40 @@ def load_model():
         logger.error(f"Failed to load model: {e}")
         return False
 
-@functools.lru_cache(maxsize=100)
 def get_card_info(hub_id: str) -> 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)
-        return "model", card.text
+        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.error(f"Error fetching model card for {hub_id}: {e}")
-        try:
-            info = dataset_info(hub_id)
-            card = DatasetCard.load(hub_id)
-            return "dataset", card.text
-        except Exception as e:
-            logger.error(f"Error fetching dataset card for {hub_id}: {e}")
-            raise ValueError(f"Could not find model or dataset with id {hub_id}")
+        logger.debug(f"No dataset card found for {hub_id}: {e}")
+
+    # Handle different cases
+    if model_exists and dataset_exists:
+        return "both", (model_text, dataset_text)
+    elif model_exists:
+        return "model", model_text
+    elif dataset_exists:
+        return "dataset", dataset_text
+    else:
+        raise ValueError(f"Could not find model or dataset with id {hub_id}")
 
 @spaces.GPU
 def _generate_summary_gpu(card_text: str, card_type: str) -> str:
@@ -71,7 +90,7 @@ def _generate_summary_gpu(card_text: str, card_type: str) -> str:
     with torch.no_grad():
         outputs = model.generate(
             inputs,
-            max_new_tokens=60,
+            max_new_tokens=0,
             pad_token_id=tokenizer.pad_token_id,
             eos_token_id=tokenizer.eos_token_id,
             temperature=0.4,
@@ -102,16 +121,29 @@ def summarize(hub_id: str = "") -> str:
         if hub_id:
             # Fetch and infer card type automatically
             card_type, card_text = get_card_info(hub_id)
-            # Use the cached wrapper
-            summary = generate_summary(card_text, card_type)
-            return f'{{"summary": "{summary}", "type": "{card_type}", "hub_id": "{hub_id}"}}'
+            
+            if card_type == "both":
+                model_text, dataset_text = card_text
+                model_summary = generate_summary(model_text, "model")
+                dataset_summary = generate_summary(dataset_text, "dataset")
+                return json.dumps({
+                    "type": "both",
+                    "hub_id": hub_id,
+                    "model_summary": model_summary,
+                    "dataset_summary": dataset_summary
+                })
+            else:
+                summary = generate_summary(card_text, card_type)
+                return json.dumps({
+                    "summary": summary,
+                    "type": card_type,
+                    "hub_id": hub_id
+                })
         else:
-            error_msg = "Error: Hub ID must be provided"
-            return f'{{"error": "{error_msg}"}}'
+            return json.dumps({"error": "Hub ID must be provided"})
 
     except Exception as e:
-        error_msg = str(e)
-        return f'{{"error": "{error_msg}"}}'
+        return json.dumps({"error": str(e)})
 
 def create_interface():
     interface = gr.Interface(
@@ -126,6 +158,6 @@ def create_interface():
 if __name__ == "__main__":
     if load_model():
         interface = create_interface()
-        interface.launch(show_error=True)
+        interface.launch()
     else:
         print("Failed to load model. Please check the logs for details.")