Update app.py

app.py

@@ -2,27 +2,25 @@ import streamlit as st
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-# Load the model and tokenizer (make sure your model is correctly loaded here)
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# Specify the model path
-model_name = "ipc_refined_approach_model"  # Replace with your actual model path or name
+# Load model and tokenizer (make sure the model path is correct)
+model_name = "ipc_refined_approach_model"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
-model = model.to(device)
 
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+# Example sections list (ensure it's aligned with your model)
+sections = ['465', '395', '332', '353', '467']
 
-# Define your legal sections
-sections = ['465', '467', '395', '332','353']  # Example sections, modify as per your actual list
+# Save labels to a file
+with open("labels.txt", "w") as f:
+    f.write("\n".join(sections))  # 'sections' should be a list of section names like ['465', '395', ...].
 
-# Streamlit UI setup
-st.title("Legal Case Section Prediction")
+# Ensure consistency in device setup
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Get input text from user
-st.subheader("Enter the legal text to predict the sections it belongs to:")
-input_text = st.text_area("Input Text", height=250)
+# Move the model to the correct device
+model = model.to(device)
 
-# Prediction function
+# Function for prediction
 def predict_text(text):
     # Tokenize and encode input text
     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
@@ -42,21 +40,23 @@ def predict_text(text):
     # Convert probabilities to binary predictions (threshold 0.5)
     predictions = {section: int(prob > 0.5) for section, prob in zip(sections, probs[0])}
     
-    # Return the sections the case belongs to
+    # Print the sections the case belongs to
     sections_belongs_to = [section for section, pred in predictions.items() if pred == 1]
-    return sections_belongs_to
-
-# Show results if input text is provided
-if input_text:
-    st.subheader("Prediction Results")
-    
-    # Get predictions for the input text
-    predicted_sections = predict_text(input_text)
-    
-    # Show predictions
-    if predicted_sections:
-        st.write(f"This case belongs to Section(s): {', '.join(predicted_sections)}")
+    if sections_belongs_to:
+        st.write(f"This case belongs to Section(s): {', '.join(sections_belongs_to)}")
     else:
         st.write("This case does not belong to any known section.")
-else:
-    st.write("Please enter some text to predict the sections.")
+    
+    return predictions
+
+# Streamlit app interface
+st.title("Legal Section Classification")
+st.write("Enter the text for case classification:")
+
+# Text input for case description
+sample_text = st.text_area("Case Description", "Attack on a police officer to avoid him from doing his duty")
+
+# Button to make prediction
+if st.button("Classify Case"):
+    predictions = predict_text(sample_text)
+    st.write(predictions)