Update app.py

app.py

@@ -1,8 +1,8 @@
 import gradio as gr
 import torch
 import torch.nn as nn
+from transformers import AutoTokenizer, AutoConfig, BertModel, BertPreTrainedModel
 from safetensors.torch import load_file
-from transformers import AutoTokenizer, AutoModel, AutoConfig
 import gc
 
 # Release memory
@@ -10,37 +10,34 @@ gc.collect()
 torch.cuda.empty_cache()
 
 model_name = "hfl/chinese-roberta-wwm-ext"
-config = AutoConfig.from_pretrained(model_name)
-
-class MultiTaskRoberta(torch.nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.bert = torch.nn.Linear(768, 768)
-        self.classifier = nn.Linear(config.hidden_size, 3)  # 3 classes for sentiment
-        self.regressor = nn.Linear(config.hidden_size, 5)   # 5 regression outputs
 
-    def forward(self, input_ids, attention_mask=None, **kwargs):
-        outputs = self.roberta(input_ids=input_ids, attention_mask=attention_mask)
-        pooled = outputs.last_hidden_state[:, 0]
-        logits = self.classifier(pooled)
-        regs = self.regressor(pooled)
-        return {"logits": logits, "regression_outputs": regs}
+class MultiTasRokBert(BertPreTrainedModel):
+    def __init__(self, config, model_name_or_path):
+        super().__init__(config)
+        # Load backbone with pretrained weights if desired
+        self.bert = BertModel.from_pretrained(model_name_or_path, config=config)
+        self.classifier = nn.Linear(config.hidden_size, 3)
+        self.regressor = nn.Linear(config.hidden_size, 5)
+
+    def forward(self, input_ids, attention_mask=None, token_type_ids=None):
+        outputs = self.bert(input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
+        pooled = outputs.pooler_output
+        sentiment_logits = self.classifier(pooled)
+        regression_outputs = self.regressor(pooled)
+        return sentiment_logits, regression_outputs
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"Device: {device}")
 
+model_path = "model1.safetensors"
+
 # Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(model_name)
+config = AutoConfig.from_pretrained(model_name)
 
-# Load base model
-base_model = AutoModel.from_pretrained(model_name)
-model = MultiTaskRoberta(config)
-
-# Load safetensors
-model_path = "model1.safetensors"
-state_dict = torch.load(model_path)
+model = MultiTaskRoBert(config, model_name).to(device)
+state_dict = load_file(model_path, device=device)
 model.load_state_dict(state_dict)
-# model.to(device)
 model.eval()
 
 # Use half precision to reduce memory usage