import torch
import torch.nn as nn
import torch.nn.functional as F
import gradio as gr
from huggingface_hub import hf_hub_download

REPO_ID  = "SentilyticsOPJ/lstm_model"
FILENAME = "lstm_model.pt"

MAX_LEN = 80
EMBEDDING_DIM = 300
LSTM_UNITS = 96
SPATIAL_DROPOUT = 0.3
DROPOUT = 0.5


class BiLSTM(nn.Module):
    def __init__(self, vocab_size, num_classes, embedding_matrix=None,
                 embedding_dim=EMBEDDING_DIM, lstm_units=LSTM_UNITS,
                 spatial_dropout=SPATIAL_DROPOUT, dropout=DROPOUT,
                 padding_idx=0):
        super().__init__()
        self.padding_idx = padding_idx

        self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=padding_idx)
        if embedding_matrix is not None:
            self.embedding.weight.data.copy_(torch.from_numpy(embedding_matrix))

        self.spatial_dropout = nn.Dropout1d(spatial_dropout)

        self.lstm = nn.LSTM(
            input_size=embedding_dim,
            hidden_size=lstm_units,
            num_layers=1,
            batch_first=True,
            bidirectional=True,
        )
        self.attn = nn.Linear(lstm_units * 2, 1)
        self.dropout = nn.Dropout(dropout)
        self.fc = nn.Linear(lstm_units * 2, num_classes)

    def forward(self, x):
        emb = self.embedding(x)                          

        emb = emb.transpose(1, 2)
        emb = self.spatial_dropout(emb)
        emb = emb.transpose(1, 2)

        out, _ = self.lstm(emb)                          

        scores = self.attn(out).squeeze(-1)              
        mask = (x == self.padding_idx)
        mask[:, 0] = False
        scores = scores.masked_fill(mask, float("-inf"))
        weights = F.softmax(scores, dim=1).unsqueeze(-1) 
        context = (out * weights).sum(dim=1)             

        return self.fc(self.dropout(context))

path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
obj  = torch.load(path, map_location="cpu", weights_only=False)
sd   = obj["state_dict"]

word2id = obj["word2id"]
labels  = obj["label_classes"]
max_len = obj["max_len"]

model = BiLSTM(len(word2id), len(labels))
model.load_state_dict(sd)
model.eval()


def predict(text):
    if not text.strip():
        return ""
    ids = [word2id.get(t, 1) for t in text.lower().split()][:max_len]
    ids += [0] * (max_len - len(ids))
    with torch.no_grad():
        probs = torch.softmax(model(torch.tensor(ids).unsqueeze(0)), dim=1)[0]
    return labels[int(probs.argmax())]


demo = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(lines=3, placeholder="Unesi tekst...", label="Text"),
    outputs=gr.Label(num_top_classes=1, label="Sentiment"),
    title="Sentiment Analysis (BiLSTM, fastText hr)",
    description="Five-class sentiment: mixed, negative, neutral, positive, sarcastic.",
    examples=["Volim kavu", "Ovaj doktor je loš", "Dan je bio ok"],
)
 
demo.launch()