Spaces:
Paused
Paused
File size: 5,299 Bytes
f006f31 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
#!/usr/bin/env python3
import pytest
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
@pytest.fixture
def classifier():
model_path = "travelgate/room_environment-classifier"
try:
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,
)
print("Modelo y tokenizer cargados exitosamente.")
except Exception as e:
print(f"Error al cargar el modelo o tokenizer: {e}")
raise e
return classification
def test_model_predictions(classifier):
test_data = [
{'input': 'comfort double with shared bed', 'expected_response': 'shared bed'},
{'input': 'habitacion con banno compartido', 'expected_response': 'shared bathroom'}, # noqa: E501
{'input': 'habitacion con bano compartido', 'expected_response': 'shared bathroom'}, # noqa: E501
{'input': 'full size bed 1 big bed with shared bathroom', 'expected_response': 'shared bathroom'}, # noqa: E501
{'input': '1 king bed suite nonsmoking shared bathroom', 'expected_response': 'shared bathroom'}, # noqa: E501
{'input': 'junior suite 1 king bed nonsmoking shared bed', 'expected_response': 'shared bed'}, # noqa: E501
{'input': 'superior sea view connected room', 'expected_response': 'connected'},
{'input': 'twin room with two shared beds', 'expected_response': 'shared bed'},
{'input': 'double room with connected rooms available', 'expected_response': 'connected'}, # noqa: E501
{'input': 'two queen beds with shared bathroom access', 'expected_response': 'shared bathroom'}, # noqa: E501
{'input': 'suite with king bed and connected room option', 'expected_response': 'connected'}, # noqa: E501
{'input': 'shared dormitory room mixed gender', 'expected_response': 'shared room'}, # noqa: E501
{'input': 'private room with shared bathroom facilities', 'expected_response': 'shared bathroom'}, # noqa: E501
{'input': 'family suite with connected rooms', 'expected_response': 'connected'}, # noqa: E501
{'input': 'bunk bed in mixed shared room', 'expected_response': 'shared room'},
{'input': 'deluxe queen room with two queen beds shared bathroom', 'expected_response': 'shared bathroom'}, # noqa: E501
{'input': 'hostel dorm bed shared room', 'expected_response': 'shared room'},
{'input': 'economy double room with shared bathroom outside the room', 'expected_response': 'shared bathroom'}, # noqa: E501
{'input': 'luxury suite with sea view and private connected rooms', 'expected_response': 'connected'}, # noqa: E501
{'input': 'single bed in co-ed shared hostel room', 'expected_response': 'shared room'}, # noqa: E501
{'input': 'executive suite with optional connecting room', 'expected_response': 'connected'}, # noqa: E501
{'input': 'standard twin room shared beds', 'expected_response': 'shared bed'},
{'input': 'cozy single room with access to shared bathroom', 'expected_response': 'shared bathroom'}, # noqa: E501
{'input': 'spacious family room with interconnecting doors', 'expected_response': 'connected'}, # noqa: E501
{'input': 'shared female dormitory for 4 guests', 'expected_response': 'shared room'}, # noqa: E501
{'input': 'luxury double room with shared bed arrangements', 'expected_response': 'shared bed'}, # noqa: E501
{'input': 'master suite with a shared bedroom balcony and sea view', 'expected_response': 'shared room'}, # noqa: E501
{'input': 'budget twin room with shared bath', 'expected_response': 'shared bathroom'}, # noqa: E501
{'input': 'two-bedroom apartment with connecting options', 'expected_response': 'connected'}, # noqa: E501
{'input': 'en-suite room with shared bathroom', 'expected_response': 'shared bathroom'}, # noqa: E501
{'input': 'single guest room with shared bed', 'expected_response': 'shared bed'}, # noqa: E501
{'input': 'shared room in a modern loft city center', 'expected_response': 'shared room'}, # noqa: E501
{'input': 'double bed in a shared room apartment', 'expected_response': 'shared room'}, # noqa: E501
{'input': 'studio with shared bath', 'expected_response': 'shared bathroom'},
{'input': 'penthouse suite connected rooms available upon request', 'expected_response': 'connected'}, # noqa: E501
{'input': "backpacker's special one bed in a shared room of six", 'expected_response': 'shared room'} # noqa: E501
]
for test_case in test_data:
description = test_case["input"]
expected_environment = test_case["expected_response"]
result = classifier(description)[0]
predicted_environment = result["label"]
assert (
predicted_environment == expected_environment
), f"Incorrect prediction for '{description}': expected '{expected_environment}', obtained '{predicted_environment}'" # noqa: E501
|