Update app.py

app.py

@@ -3,36 +3,70 @@ import streamlit as st
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import re
 import torch
+
+# Load the pre-trained model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
 model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased")
 
-def analyze_text(text):
+def analyze_text(text, confidence_threshold=0.6):
+    # Preprocess the text
     text = re.sub(r"[^\w\s]", "", text)
     text = text.lower()
+
+    # Encode the text
     encoded_text = tokenizer(text, truncation=True, padding=True, return_tensors='pt')
 
+    # Classify the text
     with torch.no_grad():
         output = model(**encoded_text)
-        predictions = output.logits.argmax(-1).item()
+        logits = output.logits
+        predictions = logits.argmax(-1).item()
+        confidence = torch.softmax(logits, dim=1)[0][predictions].item()
+
+    if confidence > confidence_threshold:
+        if predictions == 0:
+            return "Job Interview Related"
+    return "Not Job Interview Related"
+
+def count_job_related_messages(data):
+    job_related_count = 0
+    not_job_related_count = 0
+
+    for message in data["message"]:
+        result = analyze_text(message)
+        if result == "Job Interview Related":
+            job_related_count += 1
+        else:
+            not_job_related_count += 1
+
+    return job_related_count, not_job_related_count
 
-    if predictions == 0:
-        return "Job Interview Related"
-    else:
-        return "Not Job Interview Related"
+# Streamlit application
 st.title("Job Interview Message Analyzer")
 
 uploaded_file = st.file_uploader("Upload CSV file")
 user_input = st.text_input("Enter text")
 
 if uploaded_file:
+    # Read the CSV file
     data = pd.read_csv(uploaded_file)
+
+    # Analyze messages
     results = []
     for message in data["message"]:
         result = analyze_text(message)
         results.append(result)
-    data["Job_Interview_Related"] = results
+
+    data["Job Interview Related"] = results
+
+    # Count job-related messages
+    job_related_count, not_job_related_count = count_job_related_messages(data)
+
     st.dataframe(data)
+    st.write(f"Job Interview Related Messages: {job_related_count}")
+    st.write(f"Not Job Interview Related Messages: {not_job_related_count}")
 elif user_input:
+    # Analyze user-input text
     result = analyze_text(user_input)
     st.write(f"Message Classification: {result}")
 else: