Update app.py

app.py

@@ -1,27 +1,31 @@
 import streamlit as st
-from transformers import BertTokenizer, BertForSequenceClassification
 import torch
 
-BERT_PATH = "ahmedrachid/FinancialBERT"
-MODEL_PATH = "Sandy0909/finance_sentiment"
-MAX_LEN = 512
+# Model class
+class FinancialBERT(torch.nn.Module):
+    def __init__(self, model_path):
+        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)
+        return output.loss, output.logits
 
-# 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')))
+# Load model
+MODEL_PATH = "Sandy0909/finance_sentiment"
+model = FinancialBERT(MODEL_PATH)
 model.eval()
 
-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).logits
-    predictions = torch.argmax(logits, dim=-1)
-    return ['negative', 'neutral', 'positive'][predictions[0].item()]
-
-# Streamlit app
+# Streamlit App
 st.title("Financial Sentiment Analysis")
 sentence = st.text_area("Enter a financial sentence:", "")
 if st.button("Predict"):
-    sentiment = predict_sentiment(sentence)
+    # 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 = ...
+    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}")