Update models/BertTunning.py

models/BertTunning.py

@@ -1,17 +1,13 @@
-import pandas as pd
-import numpy as np
 import torch
 from torch import nn
-import torch.nn.functional as F
 
-
-from logreg_model import bert_for_logreg, tokenizer_bert
-from preprocess_bert import preprocess_bert
+from models.preprocess_stage.bert_model import model
+from models.preprocess_stage.bert_model import preprocess_bert
 
 MAX_LEN = 100
 
-class BertTunnig(nn.Module):
 
+class BertTunnig(nn.Module):
     def __init__(self, bert_model):
         super().__init__()
 
@@ -37,9 +33,9 @@ class BertTunnig(nn.Module):
         return torch.sigmoid(output)
 
 
-model_tunning = BertTunnig(bert_model=bert_for_logreg)
+model_tunning = BertTunnig(bert_model=model)
+model_tunning.load_state_dict(torch.load('models/weights/BertTunnigWeights.pt'))
 
-model_tunning.load_state_dict(torch.load('best_weights_berttinnug(2).pt'))
 
 def predict_2(text):
 
@@ -48,6 +44,6 @@ def predict_2(text):
 
     with torch.inference_mode():
 
-        predict = model_tunning(preprocessed_text, attention_mask=attention_mask).item()
+        predict = round(model_tunning(preprocessed_text, attention_mask=attention_mask).item())
 
-    return round(predict)
+    return predict