Update app.py

app.py

@@ -2,10 +2,11 @@ import torch
 import torch.nn.functional as F
 from transformers import BertTokenizer
 import gradio as gr
+import json
 
-from model import CommentMTLModel  # 导入你自定义的模型
+from model import CommentMTLModel  # 用你的类名
 
-# Set device, including MPS support for Mac
+# Set device, including MPS
 if torch.backends.mps.is_available():
     device = torch.device("mps")
 elif torch.cuda.is_available():
@@ -16,8 +17,17 @@ else:
 # Load tokenizer
 tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
 
-# Load model
-model = CommentMTLModel(config_path="config.json")
+# Load config values manually
+with open("config.json", "r") as f:
+    config_data = json.load(f)
+
+# Create model
+model = CommentMTLModel(
+    model_name="bert-base-uncased",
+    num_sentiment_labels=config_data["num_sentiment_labels"],
+    num_toxicity_labels=config_data["num_toxicity_labels"],
+    dropout_prob=config_data.get("dropout_prob", 0.1)
+)
 model.load_state_dict(torch.load("pytorch_model.bin", map_location=device))
 model.to(device)
 model.eval()
@@ -32,7 +42,10 @@ def analyse_comment(comment):
     inputs = {k: v.to(device) for k, v in inputs.items()}
     
     with torch.no_grad():
-        sentiment_logits, toxicity_logits = model(**inputs)
+        outputs = model(**inputs)
+
+    sentiment_logits = outputs["sentiment_logits"]
+    toxicity_logits = outputs["toxicity_logits"]
     
     # Process sentiment
     sentiment_probs = F.softmax(sentiment_logits, dim=1)
@@ -58,7 +71,7 @@ iface = gr.Interface(
         gr.Label(num_top_classes=1, label="Predicted Toxicity")
     ],
     title="Comment Sentiment and Toxicity Classifier",
-    description="This tool classifies the sentiment and the most probable type of toxicity in a given comment. It utilises a custom fine-tuned BERT model. Developed for academic demonstration purposes in Australia."
+    description="This tool classifies the sentiment and the most probable type of toxicity in a given comment. It utilises a custom multi-task learning BERT model. Developed for academic demonstration purposes in Australia."
 )
 
 iface.launch()