Update app.py

app.py

@@ -1,31 +1,38 @@
 import streamlit as st
+from transformers import BertTokenizer, BertForSequenceClassification
 import torch
 
-# Model class
+# Config class
+class Config:
+    TOKENIZER_PATH = "ahmedrachid/FinancialBERT"  # Use tokenizer from the original model
+    MODEL_PATH = "Sandy0909/finance_sentiment"
+    MAX_LEN = 512
+    TOKENIZER = BertTokenizer.from_pretrained(TOKENIZER_PATH)
+ 
 class FinancialBERT(torch.nn.Module):
-    def __init__(self, model_path):
+    def __init__(self):
         super(FinancialBERT, self).__init__()
-        self.bert = torch.load(model_path)
-        
-    def forward(self, input_ids, attention_mask):
-        output = self.bert(input_ids, attention_mask=attention_mask)
+        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_PATH = "Sandy0909/finance_sentiment"
-model = FinancialBERT(MODEL_PATH)
+model = FinancialBERT()
+ 
 model.eval()
-
+ 
 # Streamlit App
 st.title("Financial Sentiment Analysis")
 sentence = st.text_area("Enter a financial sentence:", "")
 if st.button("Predict"):
-    # Here, you'll need some way to tokenize the input sentence and turn it into tensors.
-    # This part has been removed in the provided code.
-    # inputs = ...
+    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)
     sentiment = ['negative', 'neutral', 'positive'][predictions[0].item()]
-    st.write(f"The predicted sentiment is: {sentiment}")
+
+    st.write(f"The predicted sentiment is: {sentiment}")