bug fix

app.py

@@ -3,11 +3,17 @@ import pandas as pd
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
 
 # Function to load the pre-trained model
-def load_model(model_name):
+def load_finetune_model(model_name):
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForSequenceClassification.from_pretrained(model_name)
     return tokenizer, model
 
+def load_model(model_name):
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSequenceClassification.from_pretrained(model_name)
+    sentiment_pipeline = pipeline("sentiment-analysis", tokenizer=tokenizer, model=model)
+    return sentiment_pipeline
+
 # Streamlit app
 st.title("Multi-label Toxicity Detection App")
 st.write("Enter a text and select the fine-tuned model to get the toxicity analysis.")
@@ -49,13 +55,9 @@ if st.button("Analyze"):
     else:
         with st.spinner("Analyzing toxicity..."):
             if selected_model == "Olivernyu/finetuned_bert_base_uncased":
-                tokenizer, model = load_model(selected_model)
-                toxicity_detector = pipeline("text-classification", tokenizer=tokenizer, model=model)
+                toxicity_detector = load_finetune_model(selected_model)
                 outputs = toxicity_detector(text, top_k=2)
-
                 category_names = ["toxic", "severe_toxic", "obscene", "threat", "insult", "identity_hate"]
-
-
                 results = []
                 for item in outputs:
                     results.append((category[item['label']], item['score']))