import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel
import gradio as gr

class MultiTaskModel(nn.Module):
    def __init__(self, base_model_name, num_topic_classes, num_sentiment_classes):
        super(MultiTaskModel, self).__init__()
        self.encoder = AutoModel.from_pretrained(base_model_name)
        hidden_size = self.encoder.config.hidden_size
        self.topik_classifier = nn.Linear(hidden_size, num_topic_classes)
        self.sentiment_classifier = nn.Linear(hidden_size, num_sentiment_classes)

    def forward(self, input_ids, attention_mask, token_type_ids=None):
        outputs = self.encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids
        )
        pooled_output = outputs.last_hidden_state[:, 0]
        topik_logits = self.topik_classifier(pooled_output)
        sentimen_logits = self.sentiment_classifier(pooled_output)
        return topik_logits, sentimen_logits

tokenizer = AutoTokenizer.from_pretrained("tokenizer")
model = MultiTaskModel("indobenchmark/indobert-base-p1", num_topic_classes=5, num_sentiment_classes=3)
model.load_state_dict(torch.load("model.pt", map_location=torch.device("cpu")))
model.eval()

topik_labels = ["Produk", "Layanan", "Pengiriman", "Pembatalan", "Lainnya"]
sentimen_labels = ["Negatif", "Netral", "Positif"]

def klasifikasi(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        topik_logits, sentimen_logits = model(**inputs)
        topik_idx = torch.argmax(topik_logits, dim=-1).item()
        sentimen_idx = torch.argmax(sentimen_logits, dim=-1).item()

    topik = topik_labels[topik_idx]
    sentimen = sentimen_labels[sentimen_idx]
    ringkasan = f"Pelanggan: {text}\nCS: Kami senang dapat membantu Anda."

    return f"HASIL ANALISIS\nTopik: {topik}\nSentimen: {sentimen}\nRingkasan: {ringkasan}"

demo = gr.Interface(fn=klasifikasi, inputs="text", outputs="text", title="Klasifikasi Topik dan Sentimen Pelanggan")
demo.launch()