Update app.py

app.py

@@ -1,79 +1,51 @@
-import gradio as gr
-from transformers import BertTokenizer, BertForSequenceClassification, BertModel, BertConfig
 import torch
-import os
+from transformers import BertTokenizer, BertModel
+import gradio as gr
 
-class MultiTaskBertModel(torch.nn.Module):
-    def __init__(self, bert_model, num_labels_task1, num_labels_task2):
+class MultiTaskBertModel(nn.Module):
+    def __init__(self, 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)
+        self.bert = BertModel.from_pretrained('bert-base-uncased')
+        self.classifier_task1 = nn.Linear(self.bert.config.hidden_size, num_labels_task1)
+        self.classifier_task2 = nn.Linear(self.bert.config.hidden_size, num_labels_task2)
+    
+    def forward(self, input_ids, attention_mask):
+        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
         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('pytorch_model.bin')
-        model.load_state_dict(torch.load(model_path))
-        
-        return model
-
-model = MultiTaskBertModel.from_pretrained("DmitriySv/ticket_classifer", 28, 3)
-tokenizer = BertTokenizer.from_pretrained("DmitriySv/ticket_classifer")
-
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model = model.to(device)
+# Загрузка сохраненной модели
+model = MultiTaskBertModel(num_labels_task1=3, num_labels_task2=4)
+model.load_state_dict(torch.load("ticket_classifier.pth"))
 model.eval()
 
-def classify(text):
-    inputs = tokenizer(text, padding=True, truncation=True, max_length=512, return_tensors="pt").to(device)
-    
-    with torch.no_grad():
-        outputs = model(**inputs)
-        print(outputs)
-        logits_task1, logits_task2 = model(**inputs)
+# Загрузка токенизатора
+tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
 
+# Функция для предсказания
+def predict(text):
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
+    input_ids = inputs["input_ids"]
+    attention_mask = inputs["attention_mask"]
+    
+    # Получаем предсказания для двух задач
+    logits_task1, logits_task2 = model(input_ids=input_ids, attention_mask=attention_mask)
+    
+    # Преобразование логитов в предсказания классов
     pred_task1 = torch.argmax(logits_task1, dim=1).item()
     pred_task2 = torch.argmax(logits_task2, dim=1).item()
-
-    return {"Тип": pred_task1, "Приоритет": pred_task2}
-
-interface = gr.Interface(
-    fn=classify,  
-    inputs=gr.Textbox(label="Введите запрос для классификации"),
-    outputs=[gr.Label(label="Тип"), gr.Label(label="Приоритет")],
-    title="Классификация запроса по типу и приоритету",
-    description="Классификация запроса по типу и приоритету."
+    
+    return {"Task 1 Prediction": pred_task1, "Task 2 Prediction": pred_task2}
+
+# Создание интерфейса с Gradio
+iface = gr.Interface(
+    fn=predict, 
+    inputs=gr.inputs.Textbox(lines=2, placeholder="Введите текст для анализа..."), 
+    outputs="json",
+    title="Multi-Task BERT Model",
+    description="Модель BERT для одновременного решения двух задач: тональность текста и тема.",
 )
 
-interface.launch()
+iface.launch()