pipeline.py

import torch
import json
import random
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel


class FoxyMultiTaskModel(nn.Module):
    """Modelo Multi-Task: Intenção + Sentimento + Urgência + Escalação Humana"""

    def __init__(self, model_name, num_intents, num_sentiments, num_urgencies):
        super().__init__()
        self.bert = AutoModel.from_pretrained(model_name)
        hidden_size = self.bert.config.hidden_size

        self.dropout = nn.Dropout(0.1)

        # Shared
        self.shared_dense = nn.Linear(hidden_size, 512)
        self.shared_activation = nn.GELU()
        self.shared_norm = nn.LayerNorm(512)

        # Intent head
        self.intent_dense1 = nn.Linear(512, 256)
        self.intent_norm1 = nn.LayerNorm(256)
        self.intent_dense2 = nn.Linear(256, 128)
        self.intent_classifier = nn.Linear(128, num_intents)

        # Sentiment head
        self.sentiment_dense = nn.Linear(512, 128)
        self.sentiment_norm = nn.LayerNorm(128)
        self.sentiment_classifier = nn.Linear(128, num_sentiments)

        # Urgency head
        self.urgency_dense = nn.Linear(512, 128)
        self.urgency_norm = nn.LayerNorm(128)
        self.urgency_classifier = nn.Linear(128, num_urgencies)

        # Human escalation head
        self.human_dense = nn.Linear(512, 64)
        self.human_classifier = nn.Linear(64, 2)

    def forward(self, input_ids=None, attention_mask=None, **kwargs):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled = self.dropout(outputs.last_hidden_state[:, 0, :])

        shared = self.shared_norm(self.shared_activation(self.shared_dense(pooled)))
        shared = self.dropout(shared)

        # Intent
        ih = self.dropout(self.intent_norm1(nn.functional.gelu(self.intent_dense1(shared))))
        ih = nn.functional.gelu(self.intent_dense2(ih))
        intent_logits = self.intent_classifier(ih)

        # Sentiment
        sh = self.dropout(self.sentiment_norm(nn.functional.gelu(self.sentiment_dense(shared))))
        sentiment_logits = self.sentiment_classifier(sh)

        # Urgency
        uh = self.dropout(self.urgency_norm(nn.functional.gelu(self.urgency_dense(shared))))
        urgency_logits = self.urgency_classifier(uh)

        # Human
        hh = nn.functional.gelu(self.human_dense(shared))
        human_logits = self.human_classifier(hh)

        return {
            "intent_logits": intent_logits,
            "sentiment_logits": sentiment_logits,
            "urgency_logits": urgency_logits,
            "human_logits": human_logits,
        }


class FoxyNLPPipeline:
    """Pipeline completo de inferência do Foxy Burger NLP."""

    def __init__(self, model_dir, device=None):
        self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")

        with open(f"{model_dir}/foxy_config.json", "r", encoding="utf-8") as f:
            self.config = json.load(f)

        self.tokenizer = AutoTokenizer.from_pretrained(model_dir)

        self.model = FoxyMultiTaskModel(
            model_name=self.config["base_model"],
            num_intents=self.config["num_intents"],
            num_sentiments=self.config["num_sentiments"],
            num_urgencies=self.config["num_urgencies"],
        )
        state_dict = torch.load(f"{model_dir}/pytorch_model.bin", map_location=self.device)
        self.model.load_state_dict(state_dict)
        self.model.to(self.device)
        self.model.eval()

        self.intent_labels = self.config["intent_labels"]
        self.sentiment_labels = self.config["sentiment_labels"]
        self.urgency_labels = self.config["urgency_labels"]
        self.responses = self.config.get("responses", {})

    def predict(self, text):
        inputs = self.tokenizer(
            text, padding="max_length", truncation=True,
            max_length=self.config["max_length"], return_tensors="pt"
        ).to(self.device)

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

        ip = torch.softmax(outputs["intent_logits"], dim=-1)[0]
        sp = torch.softmax(outputs["sentiment_logits"], dim=-1)[0]
        up = torch.softmax(outputs["urgency_logits"], dim=-1)[0]
        hp = torch.softmax(outputs["human_logits"], dim=-1)[0]

        ii, si, ui = torch.argmax(ip).item(), torch.argmax(sp).item(), torch.argmax(up).item()
        needs_human = torch.argmax(hp).item() == 1
        intent_name = self.intent_labels[ii]

        return {
            "intencao": {"label": intent_name, "confianca": round(ip[ii].item(), 4)},
            "sentimento": {"label": self.sentiment_labels[si], "confianca": round(sp[si].item(), 4)},
            "urgencia": {"label": self.urgency_labels[ui], "confianca": round(up[ui].item(), 4)},
            "precisa_humano": needs_human,
            "resposta_sugerida": random.choice(self.responses.get(intent_name, ["Vou verificar isso pra voce!"])),
            "top3_intencoes": [
                {"label": self.intent_labels[idx.item()], "confianca": round(prob.item(), 4)}
                for prob, idx in zip(*torch.topk(ip, k=3))
            ],
        }

    def predict_batch(self, texts, batch_size=32):
        results = []
        for i in range(0, len(texts), batch_size):
            batch = texts[i:i+batch_size]
            inputs = self.tokenizer(batch, padding="max_length", truncation=True,
                                    max_length=self.config["max_length"], return_tensors="pt").to(self.device)
            with torch.no_grad():
                outputs = self.model(**inputs)
            for j in range(len(batch)):
                ip = torch.softmax(outputs["intent_logits"][j], dim=-1)
                sp = torch.softmax(outputs["sentiment_logits"][j], dim=-1)
                up = torch.softmax(outputs["urgency_logits"][j], dim=-1)
                hp = torch.softmax(outputs["human_logits"][j], dim=-1)
                ii = torch.argmax(ip).item()
                intent_name = self.intent_labels[ii]
                results.append({
                    "mensagem": batch[j],
                    "intencao": intent_name,
                    "confianca": round(ip[ii].item(), 4),
                    "sentimento": self.sentiment_labels[torch.argmax(sp).item()],
                    "urgencia": self.urgency_labels[torch.argmax(up).item()],
                    "precisa_humano": torch.argmax(hp).item() == 1,
                    "resposta": random.choice(self.responses.get(intent_name, ["Vou verificar!"])),
                })
        return results


if __name__ == "__main__":
    pipe = FoxyNLPPipeline(".")
    test = pipe.predict("mano cade meu lanche faz 1 hora")
    print(json.dumps(test, indent=2, ensure_ascii=False))