use_model.py - master - danio/school_project

use_model.py
77 строк2 KB
Daniil
Initial commit
25 янв 2026, 23:05
25 янв 2026, 23:05c0de56b
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Проверка модели в консоли.
Передай токены через запятую:

python use_model.py --features "жир,шерсть"
"""

import argparse
import logging
from typing import Dict, List

import torch
import torch.nn as nn


logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(message)s",
)
logger = logging.getLogger("laundry-predict")


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:
        return self.fc(x)


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


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="model.pt")
    stain = input("Какое пятно хотите убрать? ").strip()
    item = input("Какую вещь будем стирать? ").strip()
    fabric = input("Из какой она ткани? ").strip()
    args = parser.parse_args()

    payload = torch.load(args.model, map_location="cpu")
    vocab = payload["vocab"]
    label_map = payload["label_map"]
    inv_label_map = payload["inv_label_map"]

    features = [f"fabric:{fabric}", f"item:{item}", f"stain:{stain}"]

    model = BoWClassifier(vocab_size=len(vocab), num_classes=len(label_map))
    model.load_state_dict(payload["state_dict"])
    model.eval()

    x = vectorize(features, vocab).unsqueeze(0)  # [1, vocab_size]
    with torch.no_grad():
        logits = model(x)
        probs = torch.softmax(logits, dim=1).squeeze(0)
        pred_idx = int(probs.argmax().item())
        pred_label = inv_label_map[pred_idx]
        confidence = float(probs[pred_idx].item())

    print(f"Мне кажется нужно использовать режим: {pred_label}")
    print(f"Я уверен на {int(confidence*100)}%")


if __name__ == "__main__":
    main()