Update app.py

app.py

@@ -1,48 +1,37 @@
 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 = "model.bin"
-    TRAIN_BATCH_SIZE = 32
-    VALID_BATCH_SIZE = 32
-    EPOCHS = 10
+    TOKENIZER_PATH = "ahmedrachid/FinancialBERT"  # Use tokenizer from the original model
+    MODEL_PATH = "Sandy0909/finance_sentiment"
     MAX_LEN = 512
-    TOKENIZER = BertTokenizer.from_pretrained(BERT_PATH)
+    TOKENIZER = BertTokenizer.from_pretrained(TOKENIZER_PATH)
 
 # FinancialBERT model class
-class FinancialBERT(nn.Module):
+class FinancialBERT(torch.nn.Module):
     def __init__(self):
         super(FinancialBERT, self).__init__()
-        self.bert = BertForSequenceClassification.from_pretrained(Config.BERT_PATH, num_labels=3, hidden_dropout_prob=0.5)
-        
+        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)
         return output.loss, output.logits
 
 # Load model
 model = FinancialBERT()
-model.load_state_dict(torch.load(Config.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)
+# Streamlit App
+st.title("Financial Sentiment Analysis")
+sentence = st.text_area("Enter a financial sentence:", "")
+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]
     probs = torch.nn.functional.softmax(logits, dim=-1)
     predictions = torch.argmax(probs, 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 = ['negative', 'neutral', 'positive'][predictions[0].item()]
     st.write(f"The predicted sentiment is: {sentiment}")