dataloader.py

import json
import os
from dataclasses import dataclass
from typing import List, Dict, Any, Optional, Tuple

import torch
from torch.utils.data import Dataset
from sklearn.model_selection import train_test_split
import torch.nn.functional as F  # 仅占位，如不需要可删


class JsonSentimentDataset(Dataset):
    def __init__(self, data): self.data = data
    def __len__(self): return len(self.data)
    def __getitem__(self, idx):
        item = self.data[idx]
        return {
            "reviews": item["reviews"],
            "sentiments": int(item["sentiments"]),
        }



def load_json_list(path: str) -> List[Dict[str, Any]]:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def build_datasets(
    train_path: str,
    valid_path: Optional[str] = None,
    test_path: Optional[str] = None,
    valid_size: float = 0.1,
    random_state: int = 42,
) -> Tuple[Dataset, Optional[Dataset], Optional[Dataset]]:
    train_data = load_json_list(train_path)

    if valid_path and os.path.exists(valid_path):
        valid_data = load_json_list(valid_path)
    else:
        # 从训练集中切一部分做验证（按 sentiments 分层）
        train_part, valid_part = train_test_split(
            train_data,
            test_size=valid_size,
            random_state=random_state,
            stratify=[d["sentiments"] for d in train_data],
        )
        train_data, valid_data = train_part, valid_part

    test_data = load_json_list(test_path) if test_path and os.path.exists(test_path) else None

    return (
        JsonSentimentDataset(train_data),
        JsonSentimentDataset(valid_data) if valid_data is not None else None,
        JsonSentimentDataset(test_data) if test_data is not None else None,
    )


@dataclass
class Batch:
    input_ids: torch.Tensor
    attention_mask: torch.Tensor
    labels: torch.Tensor


class DataCollatorForModernBERT:
    def __init__(self, tokenizer, max_length: int = 512):
        self.tokenizer = tokenizer
        self.max_length = max_length

    @staticmethod
    def _to_int_label(y):
        return int(y) if not isinstance(y, str) else (1 if y.strip().lower() in {"1","true","pos","positive"} else 0)

    def __call__(self, examples: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]:
        # 路径1：已 tokenized（样本含 input_ids），标签字段只用 sentiments（若已有 labels 也兼容）
        if "input_ids" in examples[0]:
            batch: Dict[str, torch.Tensor] = {}
            # 堆叠所有非标签键
            for k in examples[0].keys():
                if k in ("sentiments", "labels"):
                    continue
                vals = [e[k] for e in examples]
                batch[k] = torch.stack(vals) if isinstance(vals[0], torch.Tensor) else torch.tensor(vals)

            # 生成 labels（只从 sentiments/labels 取，不使用 label/text）
            labels = [self._to_int_label(e.get("sentiments", e.get("labels"))) for e in examples]
            batch["labels"] = torch.tensor(labels, dtype=torch.long)

            # attention_mask 兜底
            if "attention_mask" not in batch:
                pad_id = self.tokenizer.pad_token_id or 0
                batch["attention_mask"] = (batch["input_ids"] != pad_id).long()
            return batch

        # 路径2：原始 JSON（只接受 reviews/sentiments）
        texts  = [ex["reviews"] for ex in examples]
        labels = [self._to_int_label(ex["sentiments"]) for ex in examples]

        enc = self.tokenizer(
            texts, padding=True, truncation=True,
            max_length=self.max_length, return_tensors="pt",
        )
        enc["labels"] = torch.tensor(labels, dtype=torch.long)
        return enc