================================================================================
ASPEKT               | ACC      | F1-MACRO   | WYNIK
================================================================================
overall_experience   | 0.9091   | 0.7931     | (Neg/None/Pos)
location             | 0.9610   | 0.8801     | (Neg/None/Pos)
cleanliness          | 0.9655   | 0.9433     | (Neg/None/Pos)
bed_comfort          | 0.9353   | 0.8662     | (Neg/None/Pos)
size_space           | 0.8871   | 0.8226     | (Neg/None/Pos)
kitchen              | 0.9470   | 0.8923     | (Neg/None/Pos)
wifi                 | 0.9936   | 0.8931     | (Neg/None/Pos)
noise                | 0.9733   | 0.9469     | (Neg/None/Pos)
safety               | 0.9759   | 0.8595     | (Neg/None/Pos)
host_contact         | 0.9655   | 0.9415     | (Neg/None/Pos)
value_for_money      | 0.9668   | 0.8570     | (Neg/None/Pos)
food                 | 0.9428   | 0.7164     | (Neg/None/Pos)
--------------------------------------------------------------------------------

>>> SZCZEGÓŁY DLA overall_experience
              precision    recall  f1-score   support

   Negatywny       0.86      0.86      0.86       423
     Brak/OK       0.63      0.51      0.56       403
   Pozytywny       0.95      0.97      0.96      2913

    accuracy                           0.91      3739
   macro avg       0.81      0.78      0.79      3739
weighted avg       0.90      0.91      0.90      3739


>>> SZCZEGÓŁY DLA location
              precision    recall  f1-score   support

   Negatywny       0.70      0.73      0.71        92
     Brak/OK       0.96      0.95      0.95      1301
   Pozytywny       0.97      0.98      0.98      2346

    accuracy                           0.96      3739
   macro avg       0.88      0.88      0.88      3739
weighted avg       0.96      0.96      0.96      3739


>>> SZCZEGÓŁY DLA cleanliness
              precision    recall  f1-score   support

   Negatywny       0.92      0.87      0.89       302
     Brak/OK       0.97      0.98      0.98      2392
   Pozytywny       0.97      0.95      0.96      1045

    accuracy                           0.97      3739
   macro avg       0.95      0.94      0.94      3739
weighted avg       0.97      0.97      0.97      3739


>>> SZCZEGÓŁY DLA bed_comfort
              precision    recall  f1-score   support

   Negatywny       0.82      0.84      0.83       171
     Brak/OK       0.97      0.95      0.96      3102
   Pozytywny       0.76      0.87      0.81       466

    accuracy                           0.94      3739
   macro avg       0.85      0.89      0.87      3739
weighted avg       0.94      0.94      0.94      3739


>>> SZCZEGÓŁY DLA size_space
              precision    recall  f1-score   support

   Negatywny       0.77      0.83      0.80       258
     Brak/OK       0.92      0.94      0.93      2801
   Pozytywny       0.80      0.69      0.74       680

    accuracy                           0.89      3739
   macro avg       0.83      0.82      0.82      3739
weighted avg       0.88      0.89      0.89      3739


>>> SZCZEGÓŁY DLA kitchen
              precision    recall  f1-score   support

   Negatywny       0.83      0.93      0.87       164
     Brak/OK       0.98      0.96      0.97      3109
   Pozytywny       0.78      0.90      0.83       466

    accuracy                           0.95      3739
   macro avg       0.86      0.93      0.89      3739
weighted avg       0.95      0.95      0.95      3739


>>> SZCZEGÓŁY DLA wifi
              precision    recall  f1-score   support

   Negatywny       0.94      0.93      0.94        55
     Brak/OK       1.00      1.00      1.00      3648
   Pozytywny       0.72      0.78      0.75        36

    accuracy                           0.99      3739
   macro avg       0.89      0.90      0.89      3739
weighted avg       0.99      0.99      0.99      3739


>>> SZCZEGÓŁY DLA noise
              precision    recall  f1-score   support

   Negatywny       0.89      0.92      0.91       305
     Brak/OK       0.99      0.98      0.99      2881
   Pozytywny       0.95      0.95      0.95       553

    accuracy                           0.97      3739
   macro avg       0.94      0.95      0.95      3739
weighted avg       0.97      0.97      0.97      3739


>>> SZCZEGÓŁY DLA safety
              precision    recall  f1-score   support

   Negatywny       0.85      0.76      0.80       111
     Brak/OK       0.98      0.99      0.99      3490
   Pozytywny       0.87      0.72      0.79       138

    accuracy                           0.98      3739
   macro avg       0.90      0.82      0.86      3739
weighted avg       0.97      0.98      0.98      3739


>>> SZCZEGÓŁY DLA host_contact
              precision    recall  f1-score   support

   Negatywny       0.86      0.92      0.89       203
     Brak/OK       0.98      0.95      0.96      1568
   Pozytywny       0.97      0.98      0.98      1968

    accuracy                           0.97      3739
   macro avg       0.93      0.95      0.94      3739
weighted avg       0.97      0.97      0.97      3739


>>> SZCZEGÓŁY DLA value_for_money
              precision    recall  f1-score   support

   Negatywny       0.78      0.72      0.75       106
     Brak/OK       0.98      0.99      0.98      3357
   Pozytywny       0.89      0.80      0.84       276

    accuracy                           0.97      3739
   macro avg       0.88      0.83      0.86      3739
weighted avg       0.97      0.97      0.97      3739


>>> SZCZEGÓŁY DLA food
              precision    recall  f1-score   support

   Negatywny       0.48      0.58      0.52        19
     Brak/OK       0.97      0.97      0.97      3424
   Pozytywny       0.68      0.64      0.66       296

    accuracy                           0.94      3739
   macro avg       0.71      0.73      0.72      3739
weighted avg       0.94      0.94      0.94      3739

import torch
import torch.nn as nn
from typing import Dict, Optional
from transformers import XLMRobertaModel, XLMRobertaPreTrainedModel, XLMRobertaConfig, AutoTokenizer
from transformers.modeling_outputs import ModelOutput
from dataclasses import dataclass
import json


HF_REPO_ID = "bziemba/review-aspects-data-semi-synthetic-non-weighted-loss"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

ASPECT_COLUMNS = [
    "overall_experience", "location", "cleanliness", "bed_comfort",
    "size_space", "kitchen", "wifi", "noise", "safety", 
    "host_contact", "value_for_money", "food"
]

LABEL_MAP = {
    0: "negative",
    1: "none",
    2: "positive"
}

@dataclass
class MultiHeadOutput(ModelOutput):
    logits: Dict[str, torch.FloatTensor] = None

class DeepClassificationHead(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_labels, dropout_prob=0.1):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout_prob),
            nn.Linear(hidden_dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout_prob),
            nn.Linear(hidden_dim, num_labels),
        )

    def forward(self, x):
        return self.mlp(x)

class MultiHeadXLMRoberta(XLMRobertaPreTrainedModel):
    def __init__(self, config, aspect_columns=None):
        super().__init__(config)
        self.roberta = XLMRobertaModel(config)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

        self.aspect_names = aspect_columns if aspect_columns is not None else ASPECT_COLUMNS
        
        head_dim = config.hidden_size
        self.aspect_heads = nn.ModuleDict(
            {
                name: DeepClassificationHead(
                    config.hidden_size, head_dim, 3, config.hidden_dropout_prob
                )
                for name in self.aspect_names
            }
        )
        self.post_init()

    def forward(self, input_ids=None, attention_mask=None, **kwargs):
        outputs = self.roberta(input_ids, attention_mask=attention_mask)
        cls_token = self.dropout(outputs.last_hidden_state[:, 0, :])
        
        logits = {name: head(cls_token) for name, head in self.aspect_heads.items()}
        
        return MultiHeadOutput(logits=logits)



def predict(model, tokenizer, text: str, human_readable=False) -> Dict[str, str]:
    """
    Przyjmuje opinie i zwraca słownik z sentymentem dla każdego aspektu.
    Format wyjścia: {'aspect_name': 'negative' | 'none' | 'positive'}
    """
    if not text:
        return {}

    # Tokenizacja
    inputs = tokenizer(
        text, 
        return_tensors="pt", 
        truncation=True, 
        max_length=512,
        padding=True
    )
    
    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}

    with torch.no_grad():
        outputs = model(**inputs)

    results = {}
    for aspect_name, logits in outputs.logits.items():
        predicted_class_id = torch.argmax(logits, dim=-1).item()
        results[aspect_name] = LABEL_MAP[predicted_class_id]

    if human_readable:
        emojis = {
            "positive": "🟢 positive",
            "negative": "🔴 negative",
            "none": "⚪ neutral"
        }
        
        print("\n" + "-" * 40)
        for k, v in results.items():
            if v != "none" or k == 'overall_experience':
                print(f"{k:<20} : {emojis.get(v, v)}")
        print("-" * 40 + "\n")

    return results





if __name__ == "__main__":
    print(f"ładowanie modelu z {HF_REPO_ID} na {DEVICE}...")

    config = XLMRobertaConfig.from_pretrained(HF_REPO_ID)
    tokenizer = AutoTokenizer.from_pretrained(HF_REPO_ID)
    
    model = MultiHeadXLMRoberta.from_pretrained(HF_REPO_ID, config=config)
    model.to(DEVICE)
    model.eval()
    print("model załadowany.")

    sample_text = "Bardzo miła obsługa i świetna lokalizacja, ale w pokoju było trochę brudno."
    print(f"\n📝 Tekst wejściowy: {sample_text}\n")
    prediction = predict(model, tokenizer, sample_text, human_readable=True)
    # print(json.dumps(prediction, indent=4))

    sample_text = "Nasty ass room, in surprisngly nice location. Owner never showed up, didn't answered any calls."
    print(f"\n📝 Tekst wejściowy: {sample_text}\n")
    prediction = predict(model, tokenizer, sample_text, human_readable=True)
    # print(json.dumps(prediction, indent=4))

    sample_text = "Un apartamento en un barrio peligroso, lleno de burdeles y gánsteres. Wifi rapidísimo."
    print(f"\n📝 Tekst wejściowy: {sample_text}\n")
    prediction = predict(model, tokenizer, sample_text, human_readable=True)
    # print(json.dumps(prediction, indent=4))

📝 Tekst wejściowy: Bardzo miła obsługa i świetna lokalizacja, ale w pokoju było trochę brudno.
----------------------------------------
overall_experience   : ⚪ neutral
location             : 🟢 positive
cleanliness          : 🔴 negative
host_contact         : 🟢 positive
----------------------------------------


📝 Tekst wejściowy: Nasty ass room, in surprisngly nice location. Owner never showed up, didn't answered any calls.
----------------------------------------
overall_experience   : 🔴 negative
location             : 🟢 positive
cleanliness          : 🔴 negative
host_contact         : 🔴 negative
----------------------------------------


📝 Tekst wejściowy: Un apartamento en un barrio peligroso, lleno de burdeles y gánsteres. Wifi rapidísimo.
----------------------------------------
overall_experience   : 🔴 negative
location             : 🔴 negative
wifi                 : 🟢 positive
safety               : 🔴 negative
----------------------------------------