From 6e535ed0ebf9c7c55aa9dde5c8f32f29560c3553 Mon Sep 17 00:00:00 2001
From: imagineer99 <samleejackson0@gmail.com>
Date: Thu, 26 Feb 2026 06:27:52 +0000
Subject: [PATCH] fix: filter OCR datasets from non-vision hub results

---
 .../src/hooks/use-hf-dataset-search.ts        | 34 +++++++++++++++++--
 1 file changed, 32 insertions(+), 2 deletions(-)

diff --git a/studio/frontend/src/hooks/use-hf-dataset-search.ts b/studio/frontend/src/hooks/use-hf-dataset-search.ts
index b00dcee1ed..32879fd015 100644
--- a/studio/frontend/src/hooks/use-hf-dataset-search.ts
+++ b/studio/frontend/src/hooks/use-hf-dataset-search.ts
@@ -138,6 +138,7 @@ const INCOMPATIBLE_TASKS_ALL_MODELS = new Set([
 ]);
 
 const PRETRAINING_PLAIN_TAGS = new Set(["pretraining", "pre-training"]);
+const OCR_PLAIN_TAGS = new Set(["ocr", "document-ocr"]);
 
 const PRETRAINING_SIZE_CATEGORIES = new Set([
   "100M<n<1B",
@@ -147,6 +148,13 @@ const PRETRAINING_SIZE_CATEGORIES = new Set([
   "n>1T",
 ]);
 
+const OCR_OR_VISION_TEXT_TASKS = new Set([
+  "image-to-text",
+  "image-captioning",
+  "visual-question-answering",
+  "document-question-answering",
+]);
+
 const INCOMPATIBLE_TASKS_BY_MODEL: Record<ModelType, Set<string>> = {
   text: new Set([
     "text-to-image",
@@ -232,6 +240,16 @@ function rankDatasetRelevance(
 ): DatasetRelevance {
   if (isPretrainingDataset(dataset)) return "incompatible";
 
+  // Keep OCR / vision-text corpora out of non-vision defaults.
+  if (modelType !== "vision") {
+    if (
+      dataset.plainTags.some((t) => OCR_PLAIN_TAGS.has(t.toLowerCase())) ||
+      dataset.taskCategories.some((t) => OCR_OR_VISION_TEXT_TASKS.has(t))
+    ) {
+      return "incompatible";
+    }
+  }
+
   const { taskCategories } = dataset;
   if (taskCategories.length === 0) return "neutral";
 
@@ -248,6 +266,13 @@ function rankDatasetRelevance(
   return "neutral";
 }
 
+function isOcrOrVisionTextDataset(dataset: HfDatasetResult): boolean {
+  return (
+    dataset.plainTags.some((t) => OCR_PLAIN_TAGS.has(t.toLowerCase())) ||
+    dataset.taskCategories.some((t) => OCR_OR_VISION_TEXT_TASKS.has(t))
+  );
+}
+
 export function useHfDatasetSearch(
   query: string,
   options?: { modelType?: ModelType | null; accessToken?: string },
@@ -267,12 +292,17 @@ export function useHfDatasetSearch(
   const search = useHfPaginatedSearch(createIter, mapDataset);
 
   const results = useMemo(() => {
-    if (!modelType) return search.results;
+    const hideOcr = modelType !== "vision";
+    const baseResults = hideOcr
+      ? search.results.filter((ds) => !isOcrOrVisionTextDataset(ds))
+      : search.results;
+
+    if (!modelType) return baseResults;
 
     const boosted: HfDatasetResult[] = [];
     const neutral: HfDatasetResult[] = [];
 
-    for (const ds of search.results) {
+    for (const ds of baseResults) {
       const relevance = rankDatasetRelevance(ds, modelType);
       if (relevance === "boosted") boosted.push(ds);
       else if (relevance !== "incompatible") neutral.push(ds);