tests/test_roomtype.py

#!/usr/bin/env python3

import pytest
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline


@pytest.fixture
def classifier():
    model_path = "travelgate/room_type-classifier"
    model = AutoModelForSequenceClassification.from_pretrained(model_path)
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    classification = pipeline(
        "text-classification",
        model=model,
        tokenizer=tokenizer,
        framework="pt",
        device=0 if torch.cuda.is_available() else -1,
    )

    return classification


def test_model_predictions(classifier):
    test_data = [
        {'input': 'Habitacion estandar con bano', 'expected_response': 'room'},
        {'input': 'apartamento de lujo con vistas al mar', 'expected_response': 'apartment'},  # noqa: E501
        {'input': 'casa ejecutiva', 'expected_response': 'house'},
        {'input': 'villa doble', 'expected_response': 'villa'},
        {'input': 'estudio de una habitacion de lujo', 'expected_response': 'studio'},
        {'input': 'chalet premier con dos habitaciones', 'expected_response': 'chalet'},
        {'input': 'casa de la playa premium con bano compartido', 'expected_response': 'beach house'},  # noqa: E501
        {'input': 'estudio familiar grande', 'expected_response': 'family studio'},
        {'input': 'suite familiar junior', 'expected_response': 'family suite'},
        {'input': 'bungalow tradicional sin bano', 'expected_response': 'bungalow'},
        {'input': 'superior room 1 king superior room 1 king cupola or courtyard view french style 36sqm 385sq', 'expected_response': 'room'},  # noqa: E501
        {'input': 'habitacion matrimonial adaptada discapacitados', 'expected_response': 'room'},  # noqa: E501
        {'input': 'privilege room twin for 2 adults 0 children and 0 infants', 'expected_response': 'room'},  # noqa: E501
        {'input': 'deluxe room double for 2 adults 0 children and 0 infants', 'expected_response': 'room'},  # noqa: E501
        {'input': 'premier palace double room', 'expected_response': 'room'},
        {'input': 'double single use deluxe', 'expected_response': 'default'},
        {'input': 'double room queen bed superior', 'expected_response': 'room'},
        {'input': 'double guest room', 'expected_response': 'room'},
        {'input': 'single room for 1 adults 0 children and 0 infants', 'expected_response': 'room'},  # noqa: E501
        {'input': 'twin premium room incl evening tasting welcome gift comp wifi 28 sqm espresso fridge bathrobe', 'expected_response': 'room'},  # noqa: E501
        {'input': 'superior quadruple room', 'expected_response': 'room'},
        {'input': 'superior one bedroom apartment x2013 2 adults', 'expected_response': 'apartment'},  # noqa: E501
        {'input': 'deluxe room double for 2 adults 0 children and 0 infants', 'expected_response': 'room'},  # noqa: E501
        {'input': 'premier palace double room', 'expected_response': 'room'},
        {'input': 'double single use deluxe', 'expected_response': 'default'},
        {'input': 'double room queen bed superior', 'expected_response': 'room'},
        {'input': 'double guest room', 'expected_response': 'room'},
        {'input': 'single room for 1 adults 0 children and 0 infants', 'expected_response': 'room'},  # noqa: E501
        {'input': 'twin premium room incl evening tasting welcome gift comp wifi 28 sqm espresso fridge bathrobe', 'expected_response': 'room'},  # noqa: E501
        {'input': 'superior quadruple room', 'expected_response': 'room'},
        {'input': 'superior one bedroom apartment x2013 2 adults', 'expected_response': 'apartment'},  # noqa: E501
        {'input': 'comfort double', 'expected_response': 'default'},
        {'input': '1 king bed suite nonsmoking', 'expected_response': 'suite'},
        {'input': 'junior suite 1 king bed nonsmoking', 'expected_response': 'suite'},
        {'input': 'family room superior', 'expected_response': 'room'}
    ]

    for test_case in test_data:
        description = test_case["input"]
        expected_label = test_case["expected_response"]
        result = classifier(description)[0]
        predicted_label = result["label"]

        assert (
            predicted_label == expected_label
        ), f"Incorrect prediction for '{description}': expected '{expected_label}', obtained '{predicted_label}'"  # noqa: E501