Update app.py

app.py

@@ -1,8 +1,53 @@
 import gradio as gr
-from transformers import BertTokenizer, BertForSequenceClassification
+from transformers import BertTokenizer, BertForSequenceClassification, BertModel
 import torch
 
-model = BertForSequenceClassification.from_pretrained("DmitriySv/ticket_classifer")
+class MultiTaskBertModel(torch.nn.Module):
+    def __init__(self, bert_model, num_labels_task1, num_labels_task2):
+        super(MultiTaskBertModel, self).__init__()
+        self.bert = bert_model
+        self.classifier_task1 = torch.nn.Linear(self.bert.config.hidden_size, num_labels_task1)
+        self.classifier_task2 = torch.nn.Linear(self.bert.config.hidden_size, num_labels_task2)
+
+    def forward(self, input_ids, attention_mask=None, token_type_ids=None):
+        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
+        pooled_output = outputs.pooler_output
+
+        logits_task1 = self.classifier_task1(pooled_output)
+        logits_task2 = self.classifier_task2(pooled_output)
+
+        return logits_task1, logits_task2
+
+    def save_pretrained(self, save_directory):
+        # Создаем директорию, если она не существует
+        os.makedirs(save_directory, exist_ok=True)
+        
+        # Сохраняем веса модели
+        model_path = os.path.join(save_directory, 'pytorch_model.bin')
+        torch.save(self.state_dict(), model_path)
+        
+        # Сохраняем конфигурацию модели
+        config = self.bert.config
+        config.save_pretrained(save_directory)
+    
+    @classmethod
+    def from_pretrained(cls, load_directory, num_labels_task1, num_labels_task2):
+        # Загружаем конфигурацию BERT
+        config = BertConfig.from_pretrained(load_directory)
+        
+        # Загружаем BERT модель
+        bert_model = BertModel.from_pretrained(load_directory, config=config)
+        
+        # Создаем экземпляр кастомной модели
+        model = cls(bert_model, num_labels_task1, num_labels_task2)
+        
+        # Загружаем сохраненные веса
+        model_path = os.path.join(load_directory, 'pytorch_model.bin')
+        model.load_state_dict(torch.load(model_path))
+        
+        return model
+
+model = MultiTaskBertModel.from_pretrained("DmitriySv/ticket_classifer")
 tokenizer = BertTokenizer.from_pretrained("DmitriySv/ticket_classifer")
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')