Update app.py

app.py

@@ -10,44 +10,24 @@ model = BertForSequenceClassification.from_pretrained('ProsusAI/finbert')
 
 def analyze_sentiment(sec_text):
     # Encode the text
-    tokens = tokenizer.encode_plus(sec_text, add_special_tokens=False)
-    input_ids = tokens['input_ids']
-    total_len = len(input_ids)
-    attention_mask = tokens['attention_mask']
+    tokens = tokenizer.encode_plus(sec_text, add_special_tokens=True, return_tensors="pt")
     
-    # Initialize list to store probabilities
-    proba_list = []
-    
-    # Split the text into chunks for the model
-    start = 0
-    while start < total_len:
-        end = start + 510
-        if end > total_len:
-            end = total_len
-        input_ids_chunk = [101] + input_ids[start:end] + [102]  # Add [CLS] and [SEP] tokens
-        attention_mask_chunk = [1] + attention_mask[start:end] + [1]
-        with torch.no_grad():
-            outputs = model(torch.tensor([input_ids_chunk]), attention_mask=torch.tensor([attention_mask_chunk]))
-            probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)[0].tolist()
-        proba_list.append(probabilities)
-        start = end
-
-    # Calculate the mean of the probabilities
-    sentiment_scores = [sum(col) / len(col) for col in zip(*proba_list)]
-    sentiment = sentiment_scores.index(max(sentiment_scores))
+    # Make prediction
+    with torch.no_grad():
+        outputs = model(**tokens)
+        predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
 
+    # Convert predictions to sentiment labels
+    labels = ['Positive', 'Neutral', 'Negative']
+    sentiment = labels[torch.argmax(predictions)]
+    
     # Return the sentiment analysis result
-    if sentiment == 0:
-        return "Positive Sentiment"
-    elif sentiment == 1:
-        return "Negative Sentiment"
-    else:
-        return "Neutral"
+    return f"{sentiment} Sentiment"
 
 # Define the Gradio interface
 gr_interface = gr.Interface(
     fn=analyze_sentiment,
-    inputs=gr.inputs.Textbox(lines=2, placeholder="Enter SEC filing excerpt here..."),
+    inputs=gr.Textbox(lines=10, placeholder="Enter SEC filing excerpt here..."),
     outputs="text",
     title="Sentiment Analysis of SEC Filings",
     description="This tool predicts the sentiment of text excerpts from SEC filings."