Upload stance_classifier.pth

app.py

@@ -1,27 +1,41 @@
 import torch
 import torch.nn as nn
-from transformers import AutoModel, BertTokenizerFast
+from transformers import AutoTokenizer, AutoModel
 import gradio as gr
 import re
-from model import StanceClassifier
-import os
-import huggingface_hub
 
+class StanceClassifier(nn.Module):
+    def __init__(self,transformer_model, num_classes, dropout_rate=0.6):
+        super(StanceClassifier, self).__init__()
+        self.transformer = transformer_model
+        self.dropout = nn.Dropout(dropout_rate)
+        self.layer_norm = nn.LayerNorm(transformer_model.config.hidden_size)
+        self.classifier = nn.Sequential(
+            nn.Dropout(dropout_rate),
+            nn.Linear(transformer_model.config.hidden_size, transformer_model.config.hidden_size//2),
+            nn.ReLU(),
+            nn.Dropout(dropout_rate),
+            nn.Linear(transformer_model.config.hidden_size//2, num_classes)
+        )
+    def forward(self, input_ids, attention_mask):
+        outputs = self.transformer(input_ids=input_ids, attention_mask=attention_mask)
+        pooled_output = outputs.last_hidden_state[:, 0]
+        pooled_output = self.layer_norm(pooled_output)
+        logits = self.classifier(pooled_output)
+        return logits
+    
 torch.manual_seed(42)
-checkpoint = "ckiplab/bert-base-chinese"
-tokenizer = BertTokenizerFast.from_pretrained('bert-base-chinese')
+checkpoint = "bert-base-chinese"
+tokenizer = AutoTokenizer.from_pretrained(checkpoint)
 base_model = AutoModel.from_pretrained(checkpoint)
 
 model = StanceClassifier(base_model, num_classes=3)
-
-dict_path = huggingface_hub.hf_hub_download(repo_id="abcd1234davidchen/PolStanceBERT",filename="stance_classifier.pth",local_dir=".",local_dir_use_symlinks=False)
-
-model.load_state_dict(torch.load(dict_path, map_location=torch.device('cpu')))
+model.load_state_dict(torch.load("stance_classifier.pth", map_location=torch.device('cpu')))
 model.eval()
 labels = ['KMT', 'DPP', 'Neutral']
 
 def predict_stance(text):
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=64)
     with torch.no_grad():
         outputs = model(
             input_ids=inputs["input_ids"],
@@ -34,22 +48,11 @@ def predict_stance(text):
     return labels[predicted_class], confidence
 
 def gradio_interface(text):
-    singleSentenceMode = False
-    if text[0:1]=="!" or text[0:1]=="！":
-        text=text[1:]
-        singleSentenceMode = True
-
     sentences = re.split(r"[。！？\n]", text)
-    sentences = [s for idx, s in enumerate(sentences) if s.strip()]
-    accumulate_sentence = [" ".join(sentences[:idx+1]) for idx, s in enumerate(sentences) if s.strip()]
+    sentences = [s for s in sentences if s.strip()]
     results = []
-    if singleSentenceMode:
-        for s in sentences:
-            stance, conf = predict_stance(s)
-            results.append((s + f" (Confidence: {conf:.4f})", stance))
-        return results
-    for s, acus in zip(sentences, accumulate_sentence):
-        stance, conf = predict_stance(acus)
+    for s in sentences:
+        stance, conf = predict_stance(s)
         results.append((s + f" (Confidence: {conf:.4f})", stance))
     return results
 
@@ -59,7 +62,7 @@ def ui():
         inputs=gr.Textbox(label="Input Text", placeholder="Enter text to predict political stance..."),
         outputs=gr.HighlightedText(label="Prediction Result",color_map={"KMT":"blue","DPP":"green","Neutral":"purple"}),
         title="Political Stance Prediction",
-        description="Enter a text to predict its political stance (KMT, DPP, Neutral). Prefix a sentence with '!' or '！' to analyze each sentence individually.",
+        description="Enter a text to predict its political stance (KMT, DPP, Neutral)."
     ).launch()
 
 if __name__ == "__main__":

model.py

@@ -1,82 +0,0 @@
-import torch
-import torch.nn as nn
-
-
-class StanceClassifier(nn.Module):
-    def __init__(self, transformer_model, num_classes, dropout_rate=0.6):
-        super(StanceClassifier, self).__init__()
-        self.transformer = transformer_model
-        self.dropout = nn.Dropout(dropout_rate)
-        self.layer_norm = nn.LayerNorm(transformer_model.config.hidden_size)
-        
-        l0 = transformer_model.config.hidden_size
-        l1 = transformer_model.config.hidden_size * 2
-        l2 = l1 // 2
-        l3 = l2 // 2
-        # classifier expects pooled token representation (batch, hidden)
-        self.classifier = nn.Sequential(
-            nn.Linear(l0, l1),
-            nn.LayerNorm(l1),
-            nn.GELU(),
-            nn.Dropout(dropout_rate),
-            nn.Linear(l1, l2),
-            nn.LayerNorm(l2),
-            nn.GELU(),
-            nn.Dropout(dropout_rate),
-            nn.Linear(l2, l3),
-            nn.LayerNorm(l3),
-            nn.GELU(),
-            nn.Linear(l3, num_classes),
-        )
-
-        self.attention_vector = nn.Linear(l0, 1)
-        nn.init.xavier_uniform_(self.attention_vector.weight)
-
-        
-        self.freeze_transformer()
-
-    def freeze_transformer(self):
-        for param in self.transformer.parameters():
-            param.requires_grad = False
-
-    def unfreeze_transformer(self):
-        for param in self.transformer.parameters():
-            param.requires_grad = True
-
-    def forward(self, input_ids, attention_mask):
-        if not any(p.requires_grad for p in self.transformer.parameters()):
-            with torch.no_grad():
-                outputs = self.transformer(
-                    input_ids=input_ids, attention_mask=attention_mask
-                )
-        else:
-            outputs = self.transformer(
-                input_ids=input_ids, attention_mask=attention_mask
-            )
-
-        # token-level hidden states: (batch, seq_len, hidden)
-        token_states = outputs.last_hidden_state
-
-        scores = self.attention_vector(token_states).squeeze(-1)  # (batch, seq_len)
-        mask = attention_mask.to(dtype=torch.bool)  # (batch, seq_len)
-        scores = scores.masked_fill(~mask, -1e9)
-        weights = torch.softmax(scores, dim=1)  # (batch, seq_len)
-        pooled_output = (weights.unsqueeze(-1) * token_states).sum(dim=1)  # (batch, hidden)
-
-        if torch.isnan(pooled_output).any() or torch.isinf(pooled_output).any():
-            print("WARNING: Transformer output NaN/Inf")
-            pooled_output = torch.where(
-                torch.isnan(pooled_output) | torch.isinf(pooled_output),
-                torch.zeros_like(pooled_output),
-                pooled_output,
-            )
-
-        pooled_output = self.layer_norm(pooled_output)
-        logits = self.classifier(pooled_output)
-        return logits
-
-    def classifier_params(self):
-        return list(self.classifier.parameters())
-    
-    def transformer_params(self):
-        return list(self.transformer.parameters())