Update app.py

app.py

@@ -1,48 +1,27 @@
 import streamlit as st
-from transformers import BertTokenizer, BertForSequenceClassification, AdamW
+from transformers import BertTokenizer, BertForSequenceClassification
 import torch
-import torch.nn as nn
-from torch.utils.data import Dataset, DataLoader
 
-# Config class
-class Config:
-    BERT_PATH = "ahmedrachid/FinancialBERT"
-    MODEL_PATH = "Sandy0909/finance_sentiment"
-    TRAIN_BATCH_SIZE = 32
-    VALID_BATCH_SIZE = 32
-    EPOCHS = 10
-    MAX_LEN = 512
-    TOKENIZER = BertTokenizer.from_pretrained(BERT_PATH)
+BERT_PATH = "ahmedrachid/FinancialBERT"
+MODEL_PATH = "Sandy0909/finance_sentiment"
+MAX_LEN = 512
 
-# FinancialBERT model class
-class FinancialBERT(nn.Module):
-    def __init__(self):
-        super(FinancialBERT, self).__init__()
-        self.bert = BertForSequenceClassification.from_pretrained(Config.BERT_PATH, num_labels=3, hidden_dropout_prob=0.3)
-
-    def forward(self, input_ids, attention_mask, token_type_ids=None, 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.load_state_dict(torch.load(Config.MODEL_PATH, map_location=torch.device('cpu')))
+# Initialize tokenizer and model
+tokenizer = BertTokenizer.from_pretrained(BERT_PATH)
+model = BertForSequenceClassification.from_pretrained(BERT_PATH, num_labels=3)
+model.load_state_dict(torch.load(MODEL_PATH, map_location=torch.device('cpu')))
 model.eval()
 
-# Tokenizer
-tokenizer = Config.TOKENIZER
-
-def predict_sentiment(sentences):
-    inputs = tokenizer(sentences, return_tensors="pt", truncation=True, padding=True, max_length=Config.MAX_LEN)
+def predict_sentiment(sentence):
+    inputs = tokenizer(sentence, return_tensors="pt", truncation=True, padding=True, max_length=MAX_LEN)
     with torch.no_grad():
-        logits = model(**inputs)[1]
-    probs = torch.nn.functional.softmax(logits, dim=-1)
-    predictions = torch.argmax(probs, dim=-1)
+        logits = model(**inputs).logits
+    predictions = torch.argmax(logits, dim=-1)
     return ['negative', 'neutral', 'positive'][predictions[0].item()]
 
 # Streamlit app
 st.title("Financial Sentiment Analysis")
 sentence = st.text_area("Enter a financial sentence:", "")
 if st.button("Predict"):
-    sentiment = predict_sentiment([sentence])
+    sentiment = predict_sentiment(sentence)
     st.write(f"The predicted sentiment is: {sentiment}")