Update app.py

app.py

@@ -1,6 +1,6 @@
 import streamlit as st
 from transformers import pipeline
-from transformers import AutoTokenizer,AutoModelForTokenClassification,AutoModelForSequenceClassification
+from transformers import AutoTokenizer,AutoModelForTokenClassification,AutoModelForSequenceClassification,BertForSequenceClassification
 import math
 import nltk
 import torch
@@ -15,7 +15,8 @@ nltk.download('stopwords')
 
 similarityModel = SentenceTransformer('BAAI/bge-small-en')
 sentiment_model = pipeline("text-classification", model="AhmedTaha012/managersFeedback-V1.0.7")
-increase_decrease_model = pipeline("text-classification", model="AhmedTaha012/nextQuarter-status-V1.1.9")
+tokenizerQuarter = AutoTokenizer.from_pretrained('AhmedTaha012/nextQuarter-status-V1.1.9')
+modelQuarter = BertForSequenceClassification.from_pretrained('AhmedTaha012/nextQuarter-status-V1.1.9')
 tokenizerTopic = AutoTokenizer.from_pretrained("nickmuchi/finbert-tone-finetuned-finance-topic-classification",use_fast=True,token="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")
 modelTopic = AutoModelForSequenceClassification.from_pretrained("nickmuchi/finbert-tone-finetuned-finance-topic-classification",token="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")
 # torch.compile(modelTopic)
@@ -224,6 +225,13 @@ def selectedCorpusForNextQuarterModel(x,quarter,year):
     sents=[des]+courpus
     rest=[sents[f] for f in [list(cosine_similarity(embeddings_1)[0][1:]).index(value)+1 for value in sorted(list(cosine_similarity(embeddings_1)[0][1:]),reverse=True)][:3]]
     return ",".join(rest)
+def getQuarterPrediction(text):
+    tokens=tokenizerQuarter(text,padding=True,max_length=512,return_overflowing_tokens=False,add_special_tokens=True,truncation=True,return_tensors="pt")
+    with torch.no_grad():
+        logits = modelQuarter(**tokens).logits
+    predicted_class_id = logits.argmax().item()
+    return modelQuarter.config.id2label[predicted_class_id]
+
 
 st.header("Transcript Analysis", divider='rainbow')
 mainTranscript = st.text_area("Enter the transcript:", height=100)
@@ -246,8 +254,7 @@ if st.button("Analyze"):
     st.markdown(f'<span style="color:{sentiment_color}">{sentiment}</span>', unsafe_allow_html=True)
     st.subheader("Next Quarter Perdiction", divider='rainbow')
     # increase_decrease = [increase_decrease_model(x)[0]['label'] for x in chunks]
-    increase_decrease=increase_decrease_model(selectedCorpusForNextQuarterModel(mainTranscript,quarter,year))[0]['label']
-    increase_decrease=max(increase_decrease,key=increase_decrease.count)
+    increase_decrease=getQuarterPrediction(selectedCorpusForNextQuarterModel(mainTranscript,quarter,year))
     increase_decrease_color = "green" if increase_decrease == "Increase" else "red"
     st.markdown(f'<span style="color:{increase_decrease_color}">{increase_decrease}</span>', unsafe_allow_html=True)
     st.subheader("Financial Metrics", divider='rainbow')