Fix router prefix - remove duplicate /api/datasets from router definition

app/routers/datasets.py

@@ -12,7 +12,7 @@ import httpx
 from app.config import settings
 
 logger = logging.getLogger(__name__)
-router = APIRouter(prefix="/api/datasets", tags=["Datasets"])
+router = APIRouter(tags=["Datasets"])
 
 
 class DatasetInfo(BaseModel):
@@ -25,47 +25,34 @@ class DatasetInfo(BaseModel):
     likes: int = 0
     private: bool = False
     tags: List[str] = []
-    task_categories: List[str] = []
-    size_categories: List[str] = []
-    languages: List[str] = []
-    licenses: List[str] = []
     created_at: Optional[datetime]
     last_modified: Optional[datetime]
     
 
 class DatasetSearchResult(BaseModel):
-    """Search result model."""
+    """Search result dataset."""
     id: str
     author: str
     dataset_name: str
     downloads: int
     likes: int
-    description: Optional[str]
     
 
-class DatasetColumn(BaseModel):
-    """Dataset column information."""
-    name: str
-    dtype: str
-    sample_values: List[Any] = []
-    
+class DatasetRecommendation(BaseModel):
+    """Dataset recommendation for a task."""
+    dataset_id: str
+    reason: str
+    size: str
+    format: str
+
 
 class DatasetSchema(BaseModel):
     """Dataset schema information."""
+    dataset_id: str
     splits: List[str]
-    columns: List[DatasetColumn]
+    columns: List[str]
     num_rows: Dict[str, int]
-    features: Dict[str, Any]
-    
-
-class DatasetRecommendation(BaseModel):
-    """Dataset recommendation for a task."""
-    dataset_id: str
-    name: str
-    description: str
-    size: str
-    task_fit: str
-    sample_count: int
+    features: Dict[str, str]
 
 
 HF_API_URL = "https://huggingface.co/api"
@@ -74,7 +61,7 @@ HF_API_URL = "https://huggingface.co/api"
 @router.get("/search", response_model=List[DatasetSearchResult])
 async def search_datasets(
     query: str = Query(..., min_length=1, description="Search query"),
-    task: Optional[str] = Query(None, description="Filter by task category"),
+    task: Optional[str] = Query(None, description="Filter by task"),
     language: Optional[str] = Query(None, description="Filter by language"),
     author: Optional[str] = Query(None, description="Filter by author"),
     limit: int = Query(20, ge=1, le=100)
@@ -109,8 +96,7 @@ async def search_datasets(
                     author=d.get("author", ""),
                     dataset_name=d.get("id", "").split("/")[-1] if "/" in d.get("id", "") else d.get("id", ""),
                     downloads=d.get("downloads", 0),
-                    likes=d.get("likes", 0),
-                    description=d.get("description", "")[:200] if d.get("description") else None
+                    likes=d.get("likes", 0)
                 )
                 for d in data
             ]
@@ -131,27 +117,15 @@ async def get_dataset_info(dataset_id: str):
             response.raise_for_status()
             data = response.json()
             
-            # Extract metadata
-            tags = data.get("tags", [])
-            
-            task_categories = [t.replace("task_categories:", "") for t in tags if t.startswith("task_categories:")]
-            size_categories = [t.replace("size_categories:", "") for t in tags if t.startswith("size_categories:")]
-            languages = [t.replace("language:", "") for t in tags if t.startswith("language:")]
-            licenses = [t.replace("license:", "") for t in tags if t.startswith("license:")]
-            
             return DatasetInfo(
                 id=data.get("id", ""),
                 author=data.get("author", ""),
                 dataset_name=data.get("id", "").split("/")[-1] if "/" in data.get("id", "") else data.get("id", ""),
-                description=data.get("description", "")[:1000] if data.get("description") else None,
+                description=data.get("cardData", {}).get("description", ""),
                 downloads=data.get("downloads", 0),
                 likes=data.get("likes", 0),
                 private=data.get("private", False),
-                tags=tags,
-                task_categories=task_categories,
-                size_categories=size_categories,
-                languages=languages,
-                licenses=licenses,
+                tags=data.get("tags", []),
                 created_at=datetime.fromisoformat(data["createdAt"].replace("Z", "+00:00")) if data.get("createdAt") else None,
                 last_modified=datetime.fromisoformat(data["lastModified"].replace("Z", "+00:00")) if data.get("lastModified") else None
             )
@@ -164,58 +138,40 @@ async def get_dataset_info(dataset_id: str):
             raise HTTPException(status_code=500, detail=str(e))
 
 
-@router.get("/schema/{dataset_id:path}", response_model=DatasetSchema)
-async def get_dataset_schema(
-    dataset_id: str,
-    config: Optional[str] = Query(None, description="Dataset config/subset"),
-    split: str = Query("train", description="Dataset split to inspect")
-):
-    """Get schema and structure of a dataset."""
+@router.get("/schema/{dataset_id:path}")
+async def get_dataset_schema(dataset_id: str):
+    """Get dataset schema including splits, columns, and row counts."""
     try:
-        from datasets import load_dataset_builder
+        from datasets import get_dataset_infos
         
-        builder = load_dataset_builder(dataset_id, config, trust_remote_code=True)
-        info = builder.info
+        infos = get_dataset_infos(dataset_id)
+        if not infos:
+            raise HTTPException(status_code=404, detail="Dataset not found or inaccessible")
         
-        # Get splits
-        splits = list(builder.builder_configs.keys()) if builder.builder_configs else ["train"]
-        try:
-            split_dict = builder.info.splits
-            if split_dict:
-                splits = list(split_dict.keys())
-        except:
-            pass
-        
-        # Get features/columns
-        columns = []
-        features = info.features
-        if features:
-            for name, feature in features.items():
-                dtype = str(type(feature).__name__)
-                columns.append(DatasetColumn(
-                    name=name,
-                    dtype=dtype,
-                    sample_values=[]
-                ))
+        info = list(infos.values())[0]
         
-        # Get row counts
-        num_rows = {}
-        try:
-            for split_name in splits:
-                num_rows[split_name] = info.splits[split_name].num_examples if info.splits else 0
-        except:
-            num_rows = {"train": 0}
-        
-        return DatasetSchema(
-            splits=splits,
-            columns=columns,
-            num_rows=num_rows,
-            features={k: str(v) for k, v in features.items()} if features else {}
-        )
+        splits = list(info.splits.keys()) if info.splits else []
+        columns = list(info.features.keys()) if info.features else []
+        num_rows = {k: v.num_examples for k, v in info.splits.items()} if info.splits else {}
+        features = {k: str(v) for k, v in info.features.items()} if info.features else {}
         
+        return {
+            "dataset_id": dataset_id,
+            "splits": splits,
+            "columns": columns,
+            "num_rows": num_rows,
+            "features": features
+        }
     except Exception as e:
         logger.error(f"Failed to get dataset schema: {e}")
-        raise HTTPException(status_code=500, detail=f"Could not load dataset: {str(e)}")
+        return {
+            "dataset_id": dataset_id,
+            "splits": [],
+            "columns": [],
+            "num_rows": {},
+            "features": {},
+            "error": str(e)
+        }
 
 
 @router.get("/recommend/{task_type}", response_model=List[DatasetRecommendation])
@@ -223,200 +179,50 @@ async def get_dataset_recommendations(task_type: str):
     """Get dataset recommendations for a specific task."""
     recommendations = {
         "causal-lm": [
-            DatasetRecommendation(
-                dataset_id="tiiuae/falcon-refinedweb",
-                name="RefinedWeb",
-                description="High-quality web data for language modeling",
-                size="Large (500B+ tokens)",
-                task_fit="Excellent for pretraining and continued pretraining",
-                sample_count=500000000
-            ),
-            DatasetRecommendation(
-                dataset_id="OpenAssistant/oasst1",
-                name="OpenAssistant Conversations",
-                description="Human-generated assistant conversations",
-                size="Medium (84K conversations)",
-                task_fit="Instruction tuning, chat models",
-                sample_count=84000
-            ),
-            DatasetRecommendation(
-                dataset_id="HuggingFaceH4/ultrachat_200k",
-                name="UltraChat 200K",
-                description="Large-scale instruction following data",
-                size="Medium (200K examples)",
-                task_fit="Instruction tuning",
-                sample_count=200000
-            ),
-            DatasetRecommendation(
-                dataset_id="teknium/openhermes",
-                name="OpenHermes",
-                description="High-quality synthetic data",
-                size="Medium (240K examples)",
-                task_fit="Instruction tuning, chat",
-                sample_count=240000
-            ),
+            DatasetRecommendation(dataset_id="tiiuae/falcon-refinedweb", reason="Large-scale web data for language modeling", size="350GB+", format="Raw text"),
+            DatasetRecommendation(dataset_id="OpenAssistant/oasst1", reason="Human conversations for instruction tuning", size="~10M tokens", format="Conversations"),
+            DatasetRecommendation(dataset_id="HuggingFaceH4/ultrachat_200k", reason="High-quality dialogues", size="200k conversations", format="Conversations"),
+            DatasetRecommendation(dataset_id="teknium/openhermes", reason="Diverse instruction data", size="1M+ samples", format="Instructions"),
         ],
         "seq2seq": [
-            DatasetRecommendation(
-                dataset_id="cfilt/iitb-english-hindi",
-                name="IITB English-Hindi",
-                description="English-Hindi parallel corpus",
-                size="Large (1.5M pairs)",
-                task_fit="Translation",
-                sample_count=1500000
-            ),
-            DatasetRecommendation(
-                dataset_id="wmt/wmt14-en-de",
-                name="WMT14 English-German",
-                description="WMT translation benchmark",
-                size="Large (4.5M pairs)",
-                task_fit="Machine translation",
-                sample_count=4500000
-            ),
+            DatasetRecommendation(dataset_id="cfilt/iitb-english-hindi", reason="Translation dataset", size="~1.5M pairs", format="Parallel text"),
+            DatasetRecommendation(dataset_id="wmt/wmt14-en-de", reason="WMT translation benchmark", size="~4.5M pairs", format="Parallel text"),
         ],
         "summarization": [
-            DatasetRecommendation(
-                dataset_id="cnn_dailymail",
-                name="CNN/DailyMail",
-                description="News article summarization",
-                size="Large (300K articles)",
-                task_fit="News summarization",
-                sample_count=300000
-            ),
-            DatasetRecommendation(
-                dataset_id="xsum",
-                name="XSum",
-                description="BBC article summarization",
-                size="Medium (220K articles)",
-                task_fit="Abstractive summarization",
-                sample_count=220000
-            ),
-            DatasetRecommendation(
-                dataset_id="samsum",
-                name="SAMSum",
-                description="Dialogue summarization",
-                size="Small (16K dialogues)",
-                task_fit="Conversation summarization",
-                sample_count=16000
-            ),
+            DatasetRecommendation(dataset_id="cnn_dailymail", reason="News summarization standard", size="~300k articles", format="Article-Summary"),
+            DatasetRecommendation(dataset_id="xsum", reason="BBC article summaries", size="~22k articles", format="Article-Summary"),
+            DatasetRecommendation(dataset_id="samsum", reason="Dialogue summarization", size="~16k dialogues", format="Dialogue-Summary"),
         ],
         "token-classification": [
-            DatasetRecommendation(
-                dataset_id="conll2003",
-                name="CoNLL-2003",
-                description="Standard NER benchmark",
-                size="Small (22K sentences)",
-                task_fit="Named entity recognition",
-                sample_count=22000
-            ),
-            DatasetRecommendation(
-                dataset_id="wikiann",
-                name="WikiANN",
-                description="Cross-lingual NER",
-                size="Large (20M+ tags)",
-                task_fit="Multilingual NER",
-                sample_count=20000000
-            ),
-            DatasetRecommendation(
-                dataset_id="ontonotes",
-                name="OntoNotes 5.0",
-                description="Rich annotation including NER, POS",
-                size="Medium (1.8M tokens)",
-                task_fit="Comprehensive NER",
-                sample_count=1800000
-            ),
+            DatasetRecommendation(dataset_id="conll2003", reason="Classic NER benchmark", size="~22k sentences", format="IOB tags"),
+            DatasetRecommendation(dataset_id="ontonotes", reason="Multi-domain NER", size="~80k sentences", format="IOB tags"),
+            DatasetRecommendation(dataset_id="wikiann", reason="Cross-lingual NER", size="Varies by language", format="IOB tags"),
         ],
         "sequence-classification": [
-            DatasetRecommendation(
-                dataset_id="imdb",
-                name="IMDB Reviews",
-                description="Movie review sentiment",
-                size="Medium (50K reviews)",
-                task_fit="Binary sentiment classification",
-                sample_count=50000
-            ),
-            DatasetRecommendation(
-                dataset_id="yelp_review_full",
-                name="Yelp Reviews",
-                description="Business review sentiment",
-                size="Large (650K reviews)",
-                task_fit="Multi-class sentiment",
-                sample_count=650000
-            ),
-            DatasetRecommendation(
-                dataset_id="ag_news",
-                name="AG News",
-                description="News article categorization",
-                size="Medium (120K articles)",
-                task_fit="Topic classification",
-                sample_count=120000
-            ),
+            DatasetRecommendation(dataset_id="imdb", reason="Movie reviews for sentiment", size="50k reviews", format="Text-Label"),
+            DatasetRecommendation(dataset_id="yelp_polarity", reason="Yelp reviews", size="560k reviews", format="Text-Label"),
+            DatasetRecommendation(dataset_id="ag_news", reason="News topic classification", size="~127k articles", format="Text-Label"),
         ],
         "question-answering": [
-            DatasetRecommendation(
-                dataset_id="squad",
-                name="SQuAD 1.1",
-                description="Reading comprehension QA",
-                size="Medium (100K questions)",
-                task_fit="Extractive QA",
-                sample_count=100000
-            ),
-            DatasetRecommendation(
-                dataset_id="squad_v2",
-                name="SQuAD 2.0",
-                description="QA with unanswerable questions",
-                size="Large (150K questions)",
-                task_fit="Robust extractive QA",
-                sample_count=150000
-            ),
-            DatasetRecommendation(
-                dataset_id="hotpot_qa",
-                name="HotpotQA",
-                description="Multi-hop reasoning QA",
-                size="Medium (113K questions)",
-                task_fit="Multi-hop reasoning",
-                sample_count=113000
-            ),
+            DatasetRecommendation(dataset_id="squad", reason="Reading comprehension QA", size="~100k questions", format="Context-QA"),
+            DatasetRecommendation(dataset_id="squad_v2", reason="SQuAD with unanswerable", size="~150k questions", format="Context-QA"),
+            DatasetRecommendation(dataset_id="natural_questions", reason="Real user questions", size="~300k questions", format="Document-QA"),
         ],
         "translation": [
-            DatasetRecommendation(
-                dataset_id="opus100",
-                name="OPUS-100",
-                description="100 language pairs",
-                size="Very Large (100M+ pairs)",
-                task_fit="Multilingual translation",
-                sample_count=100000000
-            ),
-            DatasetRecommendation(
-                dataset_id="wmt19",
-                name="WMT19",
-                description="WMT19 translation benchmark",
-                size="Large",
-                task_fit="High-quality translation",
-                sample_count=10000000
-            ),
+            DatasetRecommendation(dataset_id="wmt/wmt14-en-de", reason="English-German", size="~4.5M pairs", format="Parallel"),
+            DatasetRecommendation(dataset_id="wmt/wmt19-ru-en", reason="Russian-English", size="~4M pairs", format="Parallel"),
         ],
         "masked-lm": [
-            DatasetRecommendation(
-                dataset_id="bookcorpus",
-                name="BookCorpus",
-                description="Books for language modeling",
-                size="Large (1B words)",
-                task_fit="Pretraining",
-                sample_count=1000000000
-            ),
-            DatasetRecommendation(
-                dataset_id="wikipedia",
-                name="Wikipedia",
-                description="Wikipedia articles",
-                size="Very Large",
-                task_fit="Pretraining, knowledge",
-                sample_count=6000000
-            ),
+            DatasetRecommendation(dataset_id="bookcorpus", reason="Books dataset for MLM", size="~1B words", format="Raw text"),
+            DatasetRecommendation(dataset_id="wikitext", reason="Wikipedia text", size="Varies", format="Raw text"),
         ],
     }
     
-    result = recommendations.get(task_type, recommendations.get("causal-lm", []))
+    result = recommendations.get(task_type, [])
+    
+    if not result:
+        result = recommendations.get("causal-lm", [])
+    
     return result[:5]
 
 
@@ -425,7 +231,7 @@ async def get_popular_datasets(
     task: Optional[str] = None,
     limit: int = Query(10, ge=1, le=50)
 ):
-    """Get most popular datasets."""
+    """Get most popular datasets, optionally filtered by task."""
     async with httpx.AsyncClient() as client:
         params = {
             "sort": "downloads",
@@ -452,72 +258,71 @@ async def get_popular_datasets(
 @router.get("/preview/{dataset_id:path}")
 async def preview_dataset(
     dataset_id: str,
-    config: Optional[str] = Query(None),
-    split: str = Query("train"),
-    limit: int = Query(5, ge=1, le=20)
+    split: str = Query("train", description="Dataset split to preview"),
+    rows: int = Query(10, ge=1, le=100)
 ):
     """Preview first N rows of a dataset."""
     try:
         from datasets import load_dataset
         
-        dataset = load_dataset(
-            dataset_id,
-            config,
-            split=split,
-            trust_remote_code=True
-        )
-        
-        rows = dataset.select(range(min(limit, len(dataset))))
+        ds = load_dataset(dataset_id, split=split, streaming=True)
+        data = []
+        for i, row in enumerate(ds):
+            if i >= rows:
+                break
+            data.append(row)
         
         return {
             "dataset_id": dataset_id,
-            "config": config,
             "split": split,
-            "total_rows": len(dataset),
-            "columns": list(dataset.features.keys()),
-            "rows": rows.to_list()
+            "rows": data,
+            "columns": list(data[0].keys()) if data else []
         }
-        
     except Exception as e:
         logger.error(f"Failed to preview dataset: {e}")
-        raise HTTPException(status_code=500, detail=f"Could not load dataset: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
 
 
 @router.get("/configs/{dataset_id:path}")
 async def get_dataset_configs(dataset_id: str):
-    """Get available configurations/subsets of a dataset."""
+    """Get available configurations for a dataset."""
     try:
         from datasets import get_dataset_config_names
         
-        configs = get_dataset_config_names(dataset_id, trust_remote_code=True)
-        
+        configs = get_dataset_config_names(dataset_id)
         return {
             "dataset_id": dataset_id,
             "configs": configs
         }
-        
     except Exception as e:
-        logger.error(f"Failed to get dataset configs: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
+        logger.error(f"Failed to get configs: {e}")
+        return {
+            "dataset_id": dataset_id,
+            "configs": [],
+            "error": str(e)
+        }
 
 
 @router.get("/splits/{dataset_id:path}")
 async def get_dataset_splits(
     dataset_id: str,
-    config: Optional[str] = Query(None)
+    config: Optional[str] = None
 ):
-    """Get available splits of a dataset."""
+    """Get available splits for a dataset."""
     try:
         from datasets import get_dataset_split_names
         
-        splits = get_dataset_split_names(dataset_id, config, trust_remote_code=True)
-        
+        splits = get_dataset_split_names(dataset_id, config_name=config)
         return {
             "dataset_id": dataset_id,
             "config": config,
-            "splits": list(splits)
+            "splits": splits
         }
-        
     except Exception as e:
-        logger.error(f"Failed to get dataset splits: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
+        logger.error(f"Failed to get splits: {e}")
+        return {
+            "dataset_id": dataset_id,
+            "config": config,
+            "splits": [],
+            "error": str(e)
+        }