Update app.py

app.py

@@ -1,7 +1,6 @@
 import streamlit as st
 import os
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
-#from transformers import GPT2Tokenizer, GPT2LMHeadModel
 import torch
 
 # Set CUDA_LAUNCH_BLOCKING for better debugging
@@ -22,12 +21,11 @@ label_mapping = {
     "Campaign": 10,
     "Intelligent Pigging": 11
 }
+
 model_name = "amiguel/names_scope_finetuned"  # Replace with your model repo
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
-#model = GPT2LMHeadModel.from_pretrained(model_name)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-#tokenizer = GPT2Tokenizer.from_pretrained(model_name)
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = torch.device("cpu")  # Force using CPU for debugging
 model.to(device)
 
 def classify_text(text, model, tokenizer, device, max_length=512):
@@ -47,7 +45,10 @@ def classify_text(text, model, tokenizer, device, max_length=512):
         logits = outputs.logits
     predicted_label = torch.argmax(logits, dim=-1).item()
     
-    return label_mapping.get(predicted_label, "Unknown")
+    if predicted_label < 0 or predicted_label >= len(label_mapping):
+        raise ValueError(f"Predicted label {predicted_label} is out of range.")
+    
+    return list(label_mapping.keys())[predicted_label]
 
 def main():
     st.title("Scope vs ItemClass")
@@ -76,7 +77,7 @@ def main():
             text_column = st.selectbox("Select the text column", df.columns)
             predicted_labels = [classify_text(text, model, tokenizer, device) for text in df[text_column]]
             df["Predicted Label"] = predicted_labels
-            st.write(df) 
+            st.write(df)
 
 if __name__ == "__main__":
     main()