Update app.py

app.py

@@ -19,7 +19,7 @@ BASELINE_PATH = os.getenv("MODEL_PATH", "baseline_pipe.pkl")
 
 
 def train_small_baseline(save_path: str = BASELINE_PATH,
-                         sample_split: str = "train[:0.5%]"):
+                         max_samples: int = 10000):
     """
     Treina um baseline pequeno usando uma amostra do dataset amazon_polarity.
     Usado apenas se baseline_pipe.pkl nao existir no Space.
@@ -27,8 +27,13 @@ def train_small_baseline(save_path: str = BASELINE_PATH,
     from datasets import load_dataset
     import pandas as pd
 
-    ds = load_dataset("amazon_polarity", split=sample_split)
-    df = pd.DataFrame({"text": ds["content"], "label": ds["label"]})
+    # Carrega o split de treino inteiro
+    ds = load_dataset("amazon_polarity", split="train")
+
+    # Embaralha e pega apenas max_samples exemplos (para ficar leve)
+    ds_small = ds.shuffle(seed=42).select(range(min(max_samples, len(ds))))
+
+    df = pd.DataFrame({"text": ds_small["content"], "label": ds_small["label"]})
 
     pipe = Pipeline(
         [
@@ -42,6 +47,7 @@ def train_small_baseline(save_path: str = BASELINE_PATH,
     return pipe
 
 
+
 def load_or_bootstrap_baseline():
     """
     Se existir baseline_pipe.pkl, carrega.