train.py - master - danio/school_project

train.py
193 строки6 KB
Daniil
Initial commit
25 янв 2026, 23:05
25 янв 2026, 23:05c0de56b
"""
Обучение простой модели: Bag-of-Words -> Линейный классификатор.
Вход: список токенов (признаки).
Выход: класс режима стирки (метка).
"""

import json
import random
import logging
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Tuple

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

# Настройка логирования (на английском языке)
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(message)s",
)
logger = logging.getLogger("laundry-train")

def read_jsonl(path: Path) -> List[dict]:
    items: List[dict] = []
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            items.append(json.loads(line))
    return items


def build_vocab(records: List[dict]) -> Dict[str, int]:
    """
    Строит словарь соответствия токенов их индексам.

    На основе всех уникальных токенов из поля "features" во всех записях.
    Токены сортируются в лексикографическом порядке для детерминированности.

    Args:
        records (List[dict]): Список записей, каждая содержит "features" — список токенов.

    Returns:
        Dict[str, int]: Словарь вида {токен: индекс}.
    """
    tokens = set()
    for r in records:
        for t in r["features"]:
            tokens.add(t)
    vocab = {tok: i for i, tok in enumerate(sorted(tokens))}
    return vocab


def build_label_map(records: List[dict]) -> Dict[str, int]:
    labels = sorted({r["label"] for r in records})
    return {lab: i for i, lab in enumerate(labels)}


def vectorize(features: List[str], vocab: Dict[str, int]) -> torch.Tensor:
    x = torch.zeros(len(vocab), dtype=torch.float32)
    for t in features:
        idx = vocab.get(t)
        if idx is not None:
            x[idx] = 1.0
    return x


@dataclass(frozen=True)
class EncodedSample:
    x: torch.Tensor  # shape: [vocab_size]
    y: int  # class index


class LaundryDataset(Dataset):
    def __init__(self, samples: List[EncodedSample]) -> None:
        self._samples = samples
    def __len__(self) -> int:
        return len(self._samples)
    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
        return s.x, torch.tensor(s.y, dtype=torch.long)


class BoWClassifier(nn.Module):
    def __init__(self, vocab_size: int, num_classes: int) -> None:
        super().__init__()
        self.fc = nn.Linear(vocab_size, num_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: [batch, vocab_size] -> logits: [batch, num_classes]
        return self.fc(x)


@torch.no_grad()
def evaluate(model: nn.Module, loader: DataLoader, device: torch.device) -> float:
    model.eval()
    correct = 0
    total = 0
    for xb, yb in loader:
        xb = xb.to(device)
        yb = yb.to(device)
        logits = model(xb)
        pred = logits.argmax(dim=1)
        correct += int((pred == yb).sum().item())
        total += int(yb.numel())
    return correct / max(total, 1)


def train(
        dataset_path: Path,
        out_path: Path,
        epochs: int = 30,
        batch_size: int = 16,
        lr: float = 0.01,
        seed: int = 42,
) -> None:
    random.seed(seed)
    torch.manual_seed(seed)

    logger.info("Loading dataset...")
    records = read_jsonl(dataset_path)

    vocab = build_vocab(records)
    label_map = build_label_map(records)
    inv_label_map = {v: k for k, v in label_map.items()}

    logger.info("Vocab size: %d", len(vocab))
    logger.info("Num classes: %d", len(label_map))

    encoded: List[EncodedSample] = []
    for r in records:
        x = vectorize(r["features"], vocab)
        y = label_map[r["label"]]
        encoded.append(EncodedSample(x=x, y=y))

    random.shuffle(encoded)
    split = int(len(encoded) * 0.8)
    train_samples = encoded[:split]
    val_samples = encoded[split:] if split < len(encoded) else encoded[:]

    train_ds = LaundryDataset(train_samples)
    val_ds = LaundryDataset(val_samples)

    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    logger.info("Using device: %s", device)

    model = BoWClassifier(vocab_size=len(vocab), num_classes=len(label_map)).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)

    logger.info("Training started...")
    for epoch in range(1, epochs + 1):
        model.train()
        total_loss = 0.0

        for xb, yb in train_loader:
            xb = xb.to(device)
            yb = yb.to(device)

            optimizer.zero_grad()
            logits = model(xb)
            loss = criterion(logits, yb)
            loss.backward()
            optimizer.step()

            total_loss += float(loss.item())

        train_acc = evaluate(model, train_loader, device)
        val_acc = evaluate(model, val_loader, device)
        avg_loss = total_loss / max(len(train_loader), 1)

        logger.info(
            "Epoch %02d/%02d | loss=%.4f | train_acc=%.2f | val_acc=%.2f",
            epoch, epochs, avg_loss, train_acc, val_acc
        )

    payload = {
        "state_dict": model.state_dict(),
        "vocab": vocab,
        "label_map": label_map,
        "inv_label_map": inv_label_map,
    }
    torch.save(payload, out_path)
    logger.info("Model saved to: %s", out_path)


if __name__ == "__main__":
    train(dataset_path=Path("dataset.jsonl"), out_path=Path("model.pt"), epochs=40, batch_size=8, lr=0.01)