abrm/multi-label-classification

1import gradio as gr2import torch3from transformers import BertTokenizer, BertModel4import pickle5import torch.nn as nn6 7# Define the model class8class CombinedAttentionTransformer(nn.Module):9    def __init__(self, num_labels, hidden_size=768, num_attention_heads=8, num_self_attention_layers=1):10        super(CombinedAttentionTransformer, self).__init__()11        self.bert = BertModel.from_pretrained('bert-base-uncased')12        self.num_labels = num_labels13        self.hidden_size = hidden_size14 15        # Self-attention layers for text sequence16        self.self_attention_layers = nn.ModuleList([17            nn.TransformerEncoderLayer(d_model=hidden_size, nhead=num_attention_heads)18            for _ in range(num_self_attention_layers)19        ])20        self.self_attention_encoder = nn.TransformerEncoder(21            nn.TransformerEncoderLayer(d_model=hidden_size, nhead=num_attention_heads),22            num_layers=num_self_attention_layers23        )24 25        # Label embeddings26        self.label_embeddings = nn.Embedding(num_labels, hidden_size)27 28        # Co-Attention layers29        self.text_to_label_attention = nn.MultiheadAttention(hidden_size, num_heads=num_attention_heads)30        self.label_to_text_attention = nn.MultiheadAttention(hidden_size, num_heads=num_attention_heads)31 32        # Fully connected layer for classification33        self.classifier = nn.Linear(hidden_size, num_labels)34        self.dropout = nn.Dropout(0.3)35 36    def forward(self, input_ids, attention_mask):37        # Get BERT output38        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)39        sequence_output = outputs[0]  # [batch_size, seq_len, hidden_size]40        pooled_output = outputs[1]  # [batch_size, hidden_size]41 42        # Apply self-attention layers to the text sequence43        for layer in self.self_attention_layers:44            sequence_output = layer(sequence_output.permute(1, 0, 2)).permute(1, 0, 2)45 46        # Create label embeddings47        label_indices = torch.arange(self.num_labels).to(input_ids.device)48        label_embeds = self.label_embeddings(label_indices).unsqueeze(1)  # [num_labels, 1, hidden_size]49 50        # Co-Attention Mechanism51        batch_size, seq_len, _ = sequence_output.size()52 53        # Step 1: Text attends to Labels54        text_attn_output, _ = self.text_to_label_attention(55            sequence_output.permute(1, 0, 2),56            label_embeds.expand(-1, batch_size, -1),57            label_embeds.expand(-1, batch_size, -1)58        )59        text_attn_output = text_attn_output.permute(1, 0, 2)  # [batch_size, seq_len, hidden_size]60 61        # Step 2: Labels attend to Text62        label_attn_output, _ = self.label_to_text_attention(63            label_embeds.expand(-1, batch_size, -1),64            sequence_output.permute(1, 0, 2),65            sequence_output.permute(1, 0, 2)66        )67        label_attn_output = label_attn_output.permute(1, 0, 2).mean(dim=1)  # [num_labels, hidden_size]68 69        # Modulate pooled output with co-attention outputs70        pooled_output = pooled_output.unsqueeze(1).repeat(1, self.num_labels, 1)  # [batch_size, num_labels, hidden_size]71        modulated_output = pooled_output * label_attn_output.unsqueeze(0).transpose(0, 1)  # [batch_size, num_labels, hidden_size]72 73        # Classification74        logits = self.classifier(self.dropout(modulated_output))  # [batch_size, num_labels, hidden_size]75        logits = logits.mean(dim=1)  # Aggregate across label dimensions76 77        return logits78 79# Load the model, tokenizer, and label columns80tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')81model = CombinedAttentionTransformer(num_labels=111)82device = torch.device('cpu')83model.load_state_dict(torch.load('MM_CombinedAttentionTransformer2.pth', map_location=device))84model.eval()85 86with open('MM_label_columns.pkl', 'rb') as f:87    label_columns = pickle.load(f)88 89# Function to predict labels for a single Input_Text90def predict_single_label(Input_Text):91    inputs = tokenizer(92        Input_Text, 93        padding='max_length', 94        truncation=True, 95        max_length=128, 96        return_tensors="pt"97    )98    input_ids = inputs['input_ids'].squeeze().unsqueeze(0).to(device)99    attention_mask = inputs['attention_mask'].squeeze().unsqueeze(0).to(device)100 101    with torch.no_grad():102        outputs = model(input_ids, attention_mask)103        predictions = torch.sigmoid(outputs).round()104    105    predicted_labels = [label_columns[i] for i, value in enumerate(predictions.cpu().numpy()[0]) if value == 1]106    return predicted_labels107 108# Function to format the output for Gradio109def format_output(labels):110    output = []111    for label in labels:112        level_1, level_2 = label.split(':')113        output.append([level_1.strip(), level_2.strip()])114    return output115 116# Gradio interface function117def gradio_predict(Input_Text):118    predicted_labels = predict_single_label(Input_Text)119    return format_output(predicted_labels)120 121# Create Gradio interface122textbox = gr.Textbox(lines=2, placeholder="Enter your text here...")123table = gr.Dataframe(headers=["Level 1", "Level 2"], datatype=["str", "str"])124 125iface = gr.Interface(126    fn=gradio_predict,127    inputs=textbox,128    outputs=table,129    title="Text Classification",130    description="Enter a text string to get predicted labels.",131    live=True132)133 134# Launch the interface135if __name__ == "__main__":136    iface.launch()