scripts/streaming_datasets_v13_9.py

"""streaming_datasets_v13_9.py — Streaming dataset loader for v13.9.1 training.

Carrega os 9 datasets especificados pelo usuário via streaming (IterableDataset),
um por vez, exaustivamente dentro do limite max_samples_per_dataset.

Datasets:
  1. CEIA-POSITIVO/ultrachat_br_clustred_balanced_v1
  2. Madras1/corpus-ptbr-v2
  3. rhaymison/multmodal_175k_portuguese
  4. TucanoBR/GigaVerbo
  5. nvidia/OpenMathReasoning
  6. MathLLMs/MathVision
  7. nvidia/OpenMathInstruct-2
  8. dominguesm/restore-punctuation-ptbr-dataset
  9. carolina-c4ai/corpus-carolina

Memory-efficient: streaming=True, no full materialization.
Returns ProcessedSample with raw_text only (tokenizer handles encoding).
"""
from __future__ import annotations

import logging
import os
import time
import traceback
from dataclasses import dataclass, field
from typing import Any, Dict, Iterator, List, Optional, Tuple

import torch

logger = logging.getLogger(__name__)

# Datasets V13.9.1 — exatamente os 9 especificados pelo usuário
DEFAULT_DATASETS = [
    "CEIA-POSITIVO/ultrachat_br_clustred_balanced_v1",
    "Madras1/corpus-ptbr-v2",
    "rhaymison/multmodal_175k_portuguese",
    "TucanoBR/GigaVerbo",
    "nvidia/OpenMathReasoning",
    "MathLLMs/MathVision",
    "nvidia/OpenMathInstruct-2",
    "dominguesm/restore-punctuation-ptbr-dataset",
    "carolina-c4ai/corpus-carolina",
]

# Dataset format descriptors — Lista priorizada de campos de texto
DATASET_FORMATS: Dict[str, Dict[str, Any]] = {
    "CEIA-POSITIVO/ultrachat_br_clustred_balanced_v1": {
        "type": "chat",
        "text_fields": ["conversa", "text", "content", "conversation", "messages"],
        "label_fields": [],
        "split": "train",
        "config": None,
    },
    "Madras1/corpus-ptbr-v2": {
        "type": "text",
        "text_fields": ["text", "content", "document", "body"],
        "label_fields": [],
        "split": "train",
        "config": None,
    },
    "rhaymison/multmodal_175k_portuguese": {
        "type": "multimodal",
        "text_fields": ["description", "text", "question", "prompt", "instruction", "input"],
        "image_fields": ["image", "image_url"],
        "label_fields": ["answer", "response", "output"],
        "split": "train",
        "config": None,
    },
    "TucanoBR/GigaVerbo": {
        "type": "text",
        "text_fields": ["text", "content", "document", "body"],
        "label_fields": [],
        "split": "train",
        "config": None,
    },
    "nvidia/OpenMathReasoning": {
        "type": "math",
        "text_fields": ["problem", "question", "input", "prompt"],
        "label_fields": ["solution", "answer", "output", "response"],
        # BUG FIX: OpenMathReasoning has splits ['cot', 'tir', 'genselect', 'additional_problems']
        "split": "cot",  # Use 'cot' (chain-of-thought) as primary split
        "config": None,
    },
    "MathLLMs/MathVision": {
        "type": "math_multimodal",
        "text_fields": ["question", "problem", "text", "query"],
        "image_fields": ["image", "image_url"],
        "label_fields": ["answer", "solution", "response"],
        # BUG FIX: MathVision has only 'test' and 'testmini' splits (no train)
        "split": "test",
        "config": None,
    },
    "nvidia/OpenMathInstruct-2": {
        "type": "math",
        "text_fields": ["problem", "question", "input", "prompt"],
        "label_fields": ["solution", "answer", "output", "response"],
        "split": "train",
        "config": None,
    },
    "dominguesm/restore-punctuation-ptbr-dataset": {
        "type": "punctuation",
        "text_fields": ["text", "original", "unpunctuated", "input"],
        "label_fields": ["punctuated", "restored", "target", "output"],
        "split": "train",
        "config": None,
    },
    "carolina-c4ai/corpus-carolina": {
        "type": "text",
        "text_fields": ["text", "content", "document", "body", "xml"],
        "label_fields": [],
        "split": "train",
        "config": None,
        # V13.9.2-carolina: dataset original usa script Python (não suportado em
        # datasets 5.0+). Workaround: stream_xml_gz direto via lxml.iterparse,
        # sem usar o script Python. Implementado em _stream_carolina_direct().
    },
    # ── V13.9.2-finetune-v2: novos datasets PT-BR ───────────────────────────
    # Adicionados para fine-tuning contínuo do modelo v13.9.2 (preserva
    # arquitetura e parâmetros). Nenhuma mudança estrutural — apenas novos
    # descritores de formato para que stream_dataset() reconheça os datasets.
    #
    # V13.9.2-finetune-v3: padronização para o formato unificado
    # "### Instruction:\n...\n\n### Response:\n..." conforme requisitado pelo
    # usuário. Apenas os datasets com estrutura instruction/response (orion,
    # cnmoro) usam o template; strak2005/bratao continuam como texto puro
    # (corpus sem estrutura de instrução). Nenhuma mudança nos 9 datasets
    # base — retrocompatibilidade total.
    "orion-research/translations-en_US-pt_BR": {
        "type": "translation",
        # string = EN (instruction), string_translation = PT-BR (response).
        # V13.9.2-finetune-v3: separa instruction/response para o template
        # unificado em vez de concatenar como texto único.
        "text_fields": ["string"],
        "label_fields": ["string_translation"],
        "split": "train",
        "config": None,
        "format_template": "instruction_response",
        "instruction_prefix": "Translate the following text to Portuguese:",
    },
    "cnmoro/Instruct-PTBR-10M": {
        "type": "instruct",
        # NB: campos em MAIÚSCULAS no schema do dataset
        # V13.9.2-finetune-v3: INSTRUCTION -> instruction, RESPONSE -> response
        # via template unificado.
        "text_fields": ["INSTRUCTION"],
        "label_fields": ["RESPONSE"],
        "split": "train",
        "config": None,
        "format_template": "instruction_response",
        # V13.9.2-finetune-v2: dataset é um único parquet de 10 GB — streaming
        # via datasets.load_dataset() é extremamente lento para materializar
        # a primeira amostra. Fallback via HF datasets-server rows API (HTTP),
        # que retorna batches de 100 rows rapidamente. Implementado em
        # _stream_via_rows_api().
        "loader": "rows_api",
    },
    "strak2005/corpus-ptbr-v1": {
        "type": "text",
        "text_fields": ["text", "content"],
        "label_fields": [],
        "split": "train",
        "config": None,
        # Fallback: se strak2005 cair, tenta bratao (mirror idêntico)
        "fallback": "bratao/corpus-ptbr-v1",
        # V13.9.2-finetune-v2: dataset tem 18 shards parquet de ~1GB cada.
        # Streaming via datasets.load_dataset() não materializa primeira
        # amostra em tempo aceitável. Fallback: baixa 1 shard localmente e
        # itera com pyarrow. Implementado em _stream_via_parquet_first_shard().
        "loader": "parquet_first_shard",
    },
    "bratao/corpus-ptbr-v1": {
        "type": "text",
        "text_fields": ["text", "content"],
        "label_fields": [],
        "split": "train",
        "config": None,
        "loader": "parquet_first_shard",
    },
}


# Lista de arquivos XML.gz do Carolina que tentaremos baixar (ordem: menores
# taxonomias primeiro para minimizar banda e RAM). Selecionamos 2 arquivos
# por taxonomia para garantir diversidade sem explodir a memória.
CAROLINA_XML_FILES = [
    # taxonomia "datasets_and_other_corpora" (DAT = pt-BR notícias)
    "corpus/datasets_and_other_corpora/pt-BR/DATa.xml.gz",
    "corpus/datasets_and_other_corpora/pt-BR/DATb.xml.gz",
    # taxonomia "wik" (wikis)
    "corpus/wik/WIKa.xml.gz",
    "corpus/wik/WIKb.xml.gz",
    # taxonomia "uni" (university_domains)
    "corpus/uni/UNIa.xml.gz",
    "corpus/uni/UNIb.xml.gz",
    # taxonomia "pub" (public_domain_works)
    "corpus/pub/PUBa.xml.gz",
    "corpus/pub/PUBb.xml.gz",
]


def _stream_carolina_direct(
    max_samples: int,
    hf_token: Optional[str] = None,
    timeout_per_file: int = 60,
) -> Iterator[Dict[str, Any]]:
    """Workaround V13.9.2-carolina: carrega Carolina sem usar o script Python.

    Baixa arquivos XML.gz diretamente do repo e faz parse streaming com
    lxml.etree.iterparse (compatível com datasets 5.0+). Cada <TEI> vira um
    documento de texto (concatenação dos <p> dentro de <body>).

    Args:
        max_samples: número máximo de amostras a retornar
        hf_token: token HF opcional (dataset é público, mas token evita rate limit)
        timeout_per_file: tempo máximo (segundos) por arquivo XML

    Yields:
        dict com chave "text" (compatível com _normalize_sample)
    """
    from huggingface_hub import hf_hub_download
    import gzip
    from lxml import etree

    TEI_NS = "{http://www.tei-c.org/ns/1.0}"
    BODY_P_TAG = f".//{TEI_NS}body/{TEI_NS}p"

    count = 0
    for xml_path in CAROLINA_XML_FILES:
        if count >= max_samples:
            return
        try:
            local_path = hf_hub_download(
                repo_id="carolina-c4ai/corpus-carolina",
                filename=xml_path,
                repo_type="dataset",
                token=hf_token,
            )
        except Exception as e:
            logger.warning(f"Carolina: não foi possível baixar {xml_path}: {e}")
            continue

        try:
            with gzip.open(local_path, "rb") as gz:
                # iterparse streaming: não carrega árvore inteira na memória
                for _, tei in etree.iterparse(
                    gz, huge_tree=True, encoding="utf-8", tag=f"{TEI_NS}TEI"
                ):
                    if count >= max_samples:
                        tei.clear()
                        return
                    # Extrai texto dos <p> dentro de <body>
                    parts = []
                    for p in tei.findall(BODY_P_TAG):
                        if p.text:
                            parts.append(p.text)
                    text = " ".join(parts).strip()
                    if text and len(text) >= 10:
                        yield {"text": text, "meta": ""}
                    tei.clear()  # libera memória da árvore TEI
                    count += 1
        except Exception as e:
            logger.warning(f"Carolina: erro processando {xml_path}: {e}")
            continue

    logger.info(f"Carolina: streamou {count} amostras no total")


# ═══════════════════════════════════════════════════════════════════════════
# V13.9.2-finetune-v2: Loaders alternativos para datasets parquet grandes
# onde datasets.load_dataset(streaming=True) é inviável (primeira amostra
# demora minutos para materializar).
# ═══════════════════════════════════════════════════════════════════════════

def _stream_via_rows_api(
    dataset_name: str,
    max_samples: int,
    hf_token: Optional[str] = None,
    batch_size: int = 100,
    max_offset: int = 100_000,
) -> Iterator[Dict[str, Any]]:
    """Carrega amostras via HF datasets-server /rows API (HTTP).

    Usado para datasets cujo parquet é muito grande para streaming eficiente
    (ex: cnmoro/Instruct-PTBR-10M com 10 GB em arquivo único). A rows API
    retorna batches de até 100 rows via HTTP, sem precisar baixar o parquet.

    Args:
        dataset_name: repo_id do dataset (ex: "cnmoro/Instruct-PTBR-10M")
        max_samples: nº máximo de amostras a retornar
        hf_token: token HF (opcional — dataset público)
        batch_size: nº de rows por requisição (máx 100)
        max_offset: offset máximo a tentar (sai do loop se excedido)

    Yields:
        dict com campos do dataset (compatível com _normalize_sample)
    """
    import json as _json
    import urllib.request
    import urllib.parse

    base_url = "https://datasets-server.huggingface.co/rows"
    count = 0
    offset = 0

    while count < max_samples and offset < max_offset:
        params = urllib.parse.urlencode({
            "dataset": dataset_name,
            "config": "default",
            "split": "train",
            "offset": offset,
            "length": min(batch_size, max_samples - count),
        })
        url = f"{base_url}?{params}"
        try:
            req = urllib.request.Request(url)
            if hf_token:
                req.add_header("Authorization", f"Bearer {hf_token}")
            with urllib.request.urlopen(req, timeout=30) as r:
                data = _json.loads(r.read())
        except Exception as e:
            logger.warning(f"rows_api {dataset_name} offset={offset}: {e}")
            break

        rows = data.get("rows", [])
        if not rows:
            break

        for row in rows:
            if count >= max_samples:
                return
            row_data = row.get("row", {})
            # Remove campos internos como __index_level_0__
            clean = {k: v for k, v in row_data.items()
                     if not k.startswith("__")}
            yield clean
            count += 1

        offset += len(rows)
        if len(rows) < batch_size:
            break  # chegou ao fim do dataset

    logger.info(f"rows_api {dataset_name}: streamou {count} amostras")


def _stream_via_parquet_first_shard(
    dataset_name: str,
    max_samples: int,
    hf_token: Optional[str] = None,
) -> Iterator[Dict[str, Any]]:
    """Baixa primeiro shard parquet do dataset e itera localmente com pyarrow.

    Usado para datasets com múltiplos shards parquet grandes (ex:
    strak2005/corpus-ptbr-v1 com 18 shards de ~1GB). Streaming via
    datasets.load_dataset() não materializa primeira amostra em tempo
    aceitável. Esta função baixa o primeiro shard (~26s para 1GB) e itera
    localmente — rápido após download.

    Args:
        dataset_name: repo_id do dataset
        max_samples: nº máximo de amostras
        hf_token: token HF opcional

    Yields:
        dict com campos do dataset
    """
    from huggingface_hub import HfApi, hf_hub_download
    import pyarrow.parquet as pq

    api = HfApi(token=hf_token)

    # Lista arquivos parquet no repo
    try:
        info = api.dataset_info(dataset_name)
    except Exception as e:
        logger.warning(f"parquet_shard {dataset_name}: dataset_info failed: {e}")
        return

    parquet_files = sorted([
        s.rfilename for s in info.siblings
        if s.rfilename.endswith(".parquet")
    ])
    if not parquet_files:
        logger.warning(f"parquet_shard {dataset_name}: nenhum .parquet encontrado")
        return

    # Baixa apenas o primeiro shard
    first_shard = parquet_files[0]
    logger.info(f"parquet_shard {dataset_name}: baixando {first_shard}...")
    try:
        local_path = hf_hub_download(
            repo_id=dataset_name,
            filename=first_shard,
            repo_type="dataset",
            token=hf_token,
        )
    except Exception as e:
        logger.warning(f"parquet_shard {dataset_name}: download failed: {e}")
        return

    # Itera com pyarrow.iter_batches — método eficiente que lê batches pequenos
    # do parquet sem materializar o row group inteiro na memória.
    # BUG FIX V13.9.2-finetune-v2: read_row_group() materializa TODAS as ~1M
    # rows de uma vez (lento + ~1GB RAM). iter_batches(batch_size=10) lê só
    # 10 rows por vez (rápido + ~10KB RAM).
    try:
        pf = pq.ParquetFile(local_path)
        count = 0
        # batch_read_size: nº de rows por batch (pequeno = menos RAM)
        for batch in pf.iter_batches(batch_size=100):
            if count >= max_samples:
                break
            cols = batch.schema.names
            # to_pylist() é eficiente para batches pequenos
            pydict = batch.to_pydict()
            n_rows = len(pydict[cols[0]]) if cols else 0
            for i in range(n_rows):
                if count >= max_samples:
                    break
                row = {col: pydict[col][i] for col in cols}
                yield row
                count += 1
    except Exception as e:
        logger.warning(f"parquet_shard {dataset_name}: pyarrow iter failed: {e}")
        return

    logger.info(f"parquet_shard {dataset_name}: streamou {count} amostras do shard {first_shard}")


@dataclass
class ProcessedSample:
    """Uma amostra processada: texto bruto + metadados."""
    raw_text: str
    dataset_name: str
    sample_idx: int
    text_field: str = ""
    label_text: str = ""


def _extract_field(sample: Dict[str, Any], candidate_fields: List[str]) -> Optional[str]:
    """Extrai texto da primeira field disponível no sample."""
    for field_name in candidate_fields:
        if field_name in sample:
            val = sample[field_name]
            if val is None:
                continue
            if isinstance(val, str):
                if val.strip():
                    return val
            elif isinstance(val, list):
                # Para chat: messages = [{role, content}, ...]
                parts = []
                for msg in val:
                    if isinstance(msg, dict):
                        role = msg.get("role", "")
                        content = msg.get("content", "")
                        if isinstance(content, str) and content.strip():
                            parts.append(f"{role}: {content}")
                    elif isinstance(msg, str) and msg.strip():
                        parts.append(msg)
                if parts:
                    return "\n".join(parts)
            elif isinstance(val, dict):
                # Pode ser nested
                text = val.get("text") or val.get("content") or val.get("body")
                if isinstance(text, str) and text.strip():
                    return text
    return None


def _normalize_sample(sample: Dict[str, Any], dataset_name: str, idx: int) -> Optional[ProcessedSample]:
    """Normaliza uma amostra bruta em ProcessedSample."""
    fmt = DATASET_FORMATS.get(dataset_name, {})
    text_fields = fmt.get("text_fields", ["text"])
    label_fields = fmt.get("label_fields", [])

    text = _extract_field(sample, text_fields)
    if not text or len(text.strip()) < 10:  # skip very short
        return None

    label = ""
    if label_fields:
        label = _extract_field(sample, label_fields) or ""

    # Se tem label, aplica template apropriado
    template = fmt.get("format_template")
    full_text = text

    # V13.9.2-finetune-v3: para o template instruction_response, exige ambos
    # os lados (instruction + response). Samples sem response são descartados
    # para evitar treinar em texto parcial.
    if template == "instruction_response" and not label:
        return None

    if label:
        # V13.9.2-finetune-v3: format_template unificado "### Instruction:/### Response:"
        # quando o dataset declarar explicitamente. Caso contrário, mantém o
        # comportamento legado (concatenação simples text\nlabel) para os 9
        # datasets base — retrocompatibilidade total.
        if template == "instruction_response":
            instr_prefix = fmt.get("instruction_prefix", "")
            if instr_prefix:
                full_text = (
                    f"### Instruction:\n{instr_prefix}\n{text}\n\n"
                    f"### Response:\n{label}"
                )
            else:
                full_text = (
                    f"### Instruction:\n{text}\n\n"
                    f"### Response:\n{label}"
                )
        else:
            # Comportamento legado (9 datasets base): concatenação simples
            full_text = f"{text}\n{label}"

    return ProcessedSample(
        raw_text=full_text,
        dataset_name=dataset_name,
        sample_idx=idx,
        text_field=text_fields[0] if text_fields else "",
        label_text=label,
    )


def _load_dataset_streaming(
    dataset_name: str,
    split: str = "train",
    config: Optional[str] = None,
    hf_token: Optional[str] = None,
):
    """Carrega dataset em modo streaming.

    BUG FIX V13.9.1: Removido trust_remote_code (deprecated).
    Se o dataset não carregar com split especificado, tenta 'train' como fallback.
    BUG FIX V13.9.2: Adicionado fallback para datasets com script (carolina-c4ai).
    """
    from datasets import load_dataset

    # Lista de splits para tentar, em ordem
    splits_to_try = [split]
    if split != "train":
        splits_to_try.append("train")
    # Adiciona outros splits comuns
    for s in ["test", "validation", "cot", "tir"]:
        if s not in splits_to_try:
            splits_to_try.append(s)

    fmt = DATASET_FORMATS.get(dataset_name, {})
    fallback_name = fmt.get("fallback")
    datasets_to_try = [dataset_name]
    if fallback_name:
        datasets_to_try.append(fallback_name)

    for try_ds_name in datasets_to_try:
        for try_split in splits_to_try:
            try:
                if config:
                    ds = load_dataset(
                        try_ds_name, config, split=try_split, streaming=True,
                        token=hf_token,
                    )
                else:
                    ds = load_dataset(
                        try_ds_name, split=try_split, streaming=True,
                        token=hf_token,
                    )
                logger.info(f"Loaded {try_ds_name} split={try_split}")
                return ds
            except Exception as e:
                err_str = str(e).lower()
                # If "Bad split" error, try next split
                if "bad split" in err_str or "available splits" in err_str:
                    logger.info(f"  Split {try_split} not available for {try_ds_name}, trying next")
                    continue
                # If "scripts no longer supported", try next dataset (fallback)
                if "scripts are no longer supported" in err_str and fallback_name:
                    logger.info(f"  {try_ds_name} uses script — trying fallback {fallback_name}")
                    break
                # For other errors, log and try next split
                logger.warning(f"Failed to load {try_ds_name} split={try_split}: {str(e)[:120]}")
                continue

    logger.warning(f"All split attempts failed for {dataset_name} (will be skipped)")
    return None


def stream_dataset(
    dataset_name: str,
    max_samples: int = 500,
    hf_token: Optional[str] = None,
) -> Iterator[ProcessedSample]:
    """Faz streaming de um dataset, retornando até max_samples amostras.

    Args:
        dataset_name: nome do dataset no HuggingFace
        max_samples: número máximo de amostras a retornar
        hf_token: token HF opcional

    Yields:
        ProcessedSample
    """
    fmt = DATASET_FORMATS.get(dataset_name)
    if fmt is None:
        logger.error(f"Unknown dataset: {dataset_name}")
        return

    # ── V13.9.2-carolina: bypass do script Python ─────────────────────────
    # O dataset carolina-c4ai/corpus-carolina usa um script Python como loader,
    # que não é suportado por datasets 5.0+. Streamamos os XML.gz diretamente.
    if dataset_name == "carolina-c4ai/corpus-carolina":
        logger.info(f"Carolina: usando carregador direto (bypass do script Python)")
        count = 0
        try:
            for raw_sample in _stream_carolina_direct(
                max_samples=max_samples, hf_token=hf_token
            ):
                processed = _normalize_sample(raw_sample, dataset_name, count)
                if processed is not None:
                    yield processed
                    count += 1
        except Exception as e:
            logger.warning(f"Carolina direct stream error at sample {count}: {e}")
            return
        logger.info(f"Streamed {count} samples from {dataset_name} (direct XML)")
        return

    # ── V13.9.2-finetune-v2: loaders alternativos para parquets grandes ───
    # Alguns datasets (cnmoro, strak2005/bratao) têm parquets tão grandes que
    # datasets.load_dataset(streaming=True) demora minutos para materializar
    # a primeira amostra. Usamos loaders alternativos especificados no campo
    # "loader" do formato.
    loader = fmt.get("loader")
    if loader == "rows_api":
        logger.info(f"{dataset_name}: usando rows_api loader (datasets-server HTTP)")
        count = 0
        try:
            for raw_sample in _stream_via_rows_api(
                dataset_name, max_samples=max_samples, hf_token=hf_token
            ):
                processed = _normalize_sample(raw_sample, dataset_name, count)
                if processed is not None:
                    yield processed
                    count += 1
        except Exception as e:
            logger.warning(f"rows_api stream error at sample {count}: {e}")
            return
        logger.info(f"Streamed {count} samples from {dataset_name} (rows_api)")
        return

    if loader == "parquet_first_shard":
        logger.info(f"{dataset_name}: usando parquet_first_shard loader")
        count = 0
        try:
            for raw_sample in _stream_via_parquet_first_shard(
                dataset_name, max_samples=max_samples, hf_token=hf_token
            ):
                processed = _normalize_sample(raw_sample, dataset_name, count)
                if processed is not None:
                    yield processed
                    count += 1
        except Exception as e:
            logger.warning(f"parquet_shard stream error at sample {count}: {e}")
            # Tenta fallback (ex: strak2005 -> bratao)
            fallback_name = fmt.get("fallback")
            if fallback_name and fallback_name != dataset_name:
                logger.info(f"Tentando fallback: {fallback_name}")
                fallback_fmt = DATASET_FORMATS.get(fallback_name, {})
                # Recursão respeitando o loader do fallback
                yield from stream_dataset(
                    fallback_name, max_samples=max_samples, hf_token=hf_token
                )
            return
        logger.info(f"Streamed {count} samples from {dataset_name} (parquet_first_shard)")
        return

    ds = _load_dataset_streaming(
        dataset_name,
        split=fmt.get("split", "train"),
        config=fmt.get("config"),
        hf_token=hf_token,
    )
    if ds is None:
        return

    count = 0
    try:
        for idx, sample in enumerate(ds):
            if count >= max_samples:
                break
            processed = _normalize_sample(sample, dataset_name, idx)
            if processed is not None:
                yield processed
                count += 1
    except Exception as e:
        logger.warning(f"Error streaming {dataset_name} at sample {idx}: {e}")
        return

    logger.info(f"Streamed {count} samples from {dataset_name}")


def stream_all_datasets(
    datasets: Optional[List[str]] = None,
    max_samples_per_dataset: int = 500,
    hf_token: Optional[str] = None,
) -> Iterator[ProcessedSample]:
    """Faz streaming sequencial de múltiplos datasets.

    Args:
        datasets: lista de nomes (default: DEFAULT_DATASETS)
        max_samples_per_dataset: limite por dataset
        hf_token: token HF opcional

    Yields:
        ProcessedSample
    """
    if datasets is None:
        datasets = DEFAULT_DATASETS

    for ds_name in datasets:
        logger.info(f"Starting dataset: {ds_name}")
        yield from stream_dataset(
            ds_name,
            max_samples=max_samples_per_dataset,
            hf_token=hf_token,
        )


def collate_samples_to_tensors(
    samples: List[ProcessedSample],
    tokenizer,
    max_seq_len: int = 64,
    pad_token_id: int = 0,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Converte lista de ProcessedSample em tensors para treino.

    BUG FIX V13.9.1: Tokenizer.encode() retorna 'Encoding' object, não list.
    Use .ids attribute to get token IDs.

    Returns:
        input_ids: (B, T)
        attention_mask: (B, T)
        labels: (B, T) — same as input_ids (LM training)
    """
    batch_input_ids = []
    batch_attn_mask = []

    for sample in samples:
        # Tokenize
        encoding = tokenizer.encode(sample.raw_text, add_special_tokens=True)

        # BUG FIX: Handle both `tokenizers.Encoding` and `transformers` tokenizer outputs
        if hasattr(encoding, 'ids'):
            tokens = encoding.ids  # tokenizers library
        elif isinstance(encoding, list):
            tokens = encoding  # already a list of ints
        elif hasattr(encoding, 'input_ids'):
            tokens = encoding.input_ids  # transformers BatchEncoding
        else:
            tokens = list(encoding)  # fallback

        # Truncate to max_seq_len
        if len(tokens) > max_seq_len:
            tokens = tokens[:max_seq_len]
        # Pad
        attn_mask = [1] * len(tokens) + [0] * (max_seq_len - len(tokens))
        tokens = list(tokens) + [pad_token_id] * (max_seq_len - len(tokens))
        batch_input_ids.append(tokens)
        batch_attn_mask.append(attn_mask)

    input_ids = torch.tensor(batch_input_ids, dtype=torch.long)
    attention_mask = torch.tensor(batch_attn_mask, dtype=torch.long)
    labels = input_ids.clone()

    return input_ids, attention_mask, labels


# ═══════════════════════════════════════════════════════════════════════════
# Self-test (NÃO enviar para HuggingFace)
# ═══════════════════════════════════════════════════════════════════════════
if __name__ == "__main__":
    print("=== Streaming Datasets V13.9.1 Self-Test ===\n")

    # Test 1: Test field extraction
    print("Test 1: Field extraction")
    sample_chat = {
        "messages": [
            {"role": "user", "content": "Olá"},
            {"role": "assistant", "content": "Como posso ajudar?"},
        ]
    }
    text = _extract_field(sample_chat, ["messages", "text"])
    print(f"  Chat text: {text!r}")
    assert text and "Olá" in text, "Field extraction failed for chat"

    sample_text = {"text": "Hello world"}
    text = _extract_field(sample_text, ["text", "content"])
    assert text == "Hello world"
    print(f"  Text field: {text!r}")
    print(f"  OK\n")

    # Test 2: Test list of datasets
    print("Test 2: Dataset list")
    print(f"  Configured datasets ({len(DEFAULT_DATASETS)}):")
    for i, ds in enumerate(DEFAULT_DATASETS, 1):
        print(f"    {i}. {ds}")
    assert len(DEFAULT_DATASETS) == 9
    print(f"  OK\n")

    print("=== ALL STREAMING TESTS PASSED ===")