Update app.py

app.py

@@ -14,13 +14,12 @@ class FinancialBERT(torch.nn.Module):
         super(FinancialBERT, self).__init__()
         self.bert = BertForSequenceClassification.from_pretrained(Config.MODEL_PATH, num_labels=3, hidden_dropout_prob=0.5)
  
-    def forward(self, input_ids, attention_mask, labels=None):
-        output = self.bert(input_ids, attention_mask=attention_mask, labels=labels)
+    def forward(self, input_ids, attention_mask, token_type_ids, labels=None):
+        output = self.bert(input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, labels=labels)
         return output.loss, output.logits
  
 # Load model
 model = FinancialBERT()
- 
 model.eval()
  
 # Streamlit App
@@ -30,9 +29,9 @@ if st.button("Predict"):
     tokenizer = Config.TOKENIZER
     inputs = tokenizer([sentence], return_tensors="pt", truncation=True, padding=True, max_length=Config.MAX_LEN)
     with torch.no_grad():
-        logits = model(**inputs)[1]
+        logits = model(input_ids=inputs['input_ids'], attention_mask=inputs['attention_mask'], token_type_ids=inputs.get('token_type_ids'))[1]
     probs = torch.nn.functional.softmax(logits, dim=-1)
     predictions = torch.argmax(probs, dim=-1)
     sentiment = ['negative', 'neutral', 'positive'][predictions[0].item()]
 
-    st.write(f"The predicted sentiment is: {sentiment}")
+    st.write(f"The predicted sentiment is: {sentiment}")