yihan1/restaurant-review-analyzer-mt5

Restaurant Review Analyzer

This model analyzes Dutch restaurant reviews and predicts scores across three dimensions:

• Taste
• Service
• Ambiance

Model Details

• Model Type: Multi-head regression model
• Language: Multi-lingual
• Base Model: google/mt5-base
• Training Dataset: NL_restaurant_reviews
• Output: Scores on a 1-10 scale for each dimension

Model Evaluation

The model was evaluated on the tokenized_test set using the following regression metrics:

Dimension	MSE	MAE	R²
Taste	1.2105	0.8216	0.7334
Service	1.4076	0.8862	0.7229
Ambiance	1.4333	0.9073	0.4796
Overall	1.3505	0.8717	0.6453

These results show that the model is most accurate in predicting Taste and Service, with somewhat lower performance on Ambiance, likely due to higher subjectivity in those ratings.

Note: Metrics were calculated by comparing model predictions with the ground-truth scores on a held-out test set using standard regression evaluation metrics.

Usage

import os
import torch
import torch.nn as nn
import numpy as np
import pandas as pd
import json
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader, Dataset
from transformers import MT5Model, MT5Tokenizer
from transformers import get_cosine_schedule_with_warmup
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from datasets import load_dataset, Dataset
from tqdm.auto import tqdm

# Load the model
class RestaurantReviewAnalyzer(nn.Module):
    def __init__(self, pretrained_model_name="google/mt5-base", num_dimensions=3):
        super().__init__()
        # Load pre-trained encoder
        self.encoder = MT5Model.from_pretrained(pretrained_model_name)

        # Get encoder output dimension
        hidden_size = self.encoder.encoder.config.hidden_size

        # Define dimension names
        self.dimension_names = ["Taste", "Service", "Ambiance"]

        # Create regression heads for each dimension with improved architecture
        self.regression_heads = nn.ModuleDict({
            dim: nn.Sequential(
                nn.Dropout(0.2),
                nn.Linear(hidden_size, 256),
                nn.BatchNorm1d(256),
                nn.GELU(),
                nn.Dropout(0.1),
                nn.Linear(256, 128),
                nn.BatchNorm1d(128),
                nn.GELU(),
                nn.Linear(128, 1)
            ) for dim in self.dimension_names[:num_dimensions]
        })

    def forward(self, input_ids=None, attention_mask=None):
        # Get encoding from the base model
        encoder_output = self.encoder.encoder(
            input_ids=input_ids,
            attention_mask=attention_mask
        )

        # Use [CLS] token as the review representation
        sequence_output = encoder_output.last_hidden_state[:, 0, :]

        # Apply regression heads to predict scores for each dimension
        results = {}
        for dim_name, regression_head in self.regression_heads.items():
            # Scale to 1-10 range
            score = regression_head(sequence_output)
            results[dim_name] = 1.0 + 9.0 * torch.sigmoid(score)

        return results

# Load model and tokenizer
model_path = "yihan1/restaurant-review-analyzer-mt5"
tokenizer = MT5Tokenizer.from_pretrained(model_path)

# Load model config
with open(os.path.join(model_path, "model_config.json"), "r") as f:
    model_config = json.load(f)

# Initialize model with encoder
model = RestaurantReviewAnalyzer(pretrained_model_name=model_path)

# Load regression heads
with open(os.path.join(model_path, "regression_heads.json"), "r") as f:
    regression_heads_dict = json.load(f)

# Convert regression head parameters back to tensors
for dim_name, params in regression_heads_dict.items():
    for k, v in params.items():
        model.regression_heads[dim_name].state_dict()[k].copy_(torch.tensor(v))

# Example inference
review = "Heerlijk gegeten bij dit restaurant! De service was top en de sfeer gezellig."
inputs = tokenizer(review, return_tensors="pt", padding=True, truncation=True, max_length=512)
with torch.no_grad():
    outputs = model(inputs["input_ids"], inputs["attention_mask"])

for dim, score in outputs.items():
    print(f"{dim}: {score.item():.1f}/10")