Update app.py

app.py

@@ -1,24 +1,44 @@
 #%%writefile debias_app.py
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModelForTokenClassification, pipeline as tf_pipeline
 import transformers
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModelForTokenClassification, pipeline as tf_pipeline
 import torch
 import pandas as pd
 
-
 # Set the default device
+torch_device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+# Model Loader Class for Lazy Loading and Resource Sharing
+class ModelLoader:
+    def __init__(self):
+        self.tokenizers = {}
+        self.models = {}
+
+    def load_tokenizer(self, model_name):
+        if model_name not in self.tokenizers:
+            self.tokenizers[model_name] = AutoTokenizer.from_pretrained(model_name)
+        return self.tokenizers[model_name]
 
+    def load_model(self, model_type, model_name):
+        if model_name not in self.models:
+            if model_type == "classification":
+                self.models[model_name] = AutoModelForSequenceClassification.from_pretrained(model_name).to(torch_device)
+            elif model_type == "token_classification":
+                self.models[model_name] = AutoModelForTokenClassification.from_pretrained(model_name).to(torch_device)
+        return self.models[model_name]
 
+model_loader = ModelLoader()
+
+# BiasPipeline for handling bias detection and cleaning
 class BiasPipeline:
-    def __init__(self):
-        # Load models and tokenizers for bias detection
-        self.classifier_tokenizer = AutoTokenizer.from_pretrained("newsmediabias/UnBIAS-classification-bert")
-        self.classifier_model = AutoModelForSequenceClassification.from_pretrained("newsmediabias/UnBIAS-classification-bert")
-        self.ner_tokenizer = AutoTokenizer.from_pretrained("newsmediabias/UnBIAS-NER")
-        self.ner_model = AutoModelForTokenClassification.from_pretrained("newsmediabias/UnBIAS-NER")
-        self.classifier = tf_pipeline("text-classification", model=self.classifier_model, tokenizer=self.classifier_tokenizer)
-        self.ner = tf_pipeline("ner", model=self.ner_model, tokenizer=self.ner_tokenizer)
- 
+    def __init__(self, model_loader):
+        self.model_loader = model_loader
+        self.classifier_tokenizer = model_loader.load_tokenizer("newsmediabias/UnBIAS-classification-bert")
+        self.classifier_model = model_loader.load_model("classification", "newsmediabias/UnBIAS-classification-bert")
+        self.ner_tokenizer = model_loader.load_tokenizer("newsmediabias/UnBIAS-NER")
+        self.ner_model = model_loader.load_model("token_classification", "newsmediabias/UnBIAS-NER")
+        self.classifier = tf_pipeline("text-classification", model=self.classifier_model, tokenizer=self.classifier_tokenizer, device=0)
+        self.ner = tf_pipeline("ner", model=self.ner_model, tokenizer=self.ner_tokenizer, device=0)
 
     def clean_text(self, text):
         """Clean up the text by removing any redundant spaces."""
@@ -30,72 +50,11 @@ class BiasPipeline:
         ner_results = self.ner(texts)
         return classification_results, ner_results
 
-
-# Constants and Global Variables
-sys_message = "Task: Please just generate a bias-free version of the text provided, ensuring it's free from biases related to  age, gender, politics, social nuances, or economic background, while keeping it roughly the same length as the original:"
-instruction = "Instruction: As a helpful, respectful and trustworthy debiasing assistant, your task is to receive a text and return its unbiased version, Don't add additional comment. Just return the  un biased version of the input text:"
-# Model setup for debiasing
-debias_model = "newsmediabias/UnBIAS-LLama2-Debiaser-Chat-QLoRA"
-debias_tokenizer = AutoTokenizer.from_pretrained(debias_model)
-debias_pipeline = transformers.pipeline(
-    "text-generation",
-    model=debias_model,
-    torch_dtype=torch.float16,
-    device_map="auto",
-)
-
-
-
-def get_debiased_sequence(prompt):
-    """
-    Generate a debiased version of the provided text using the debiasing pipeline.
-    Args:
-    - prompt (str): Text to be debiased.
-    Returns:
-    - str: Debiased text.
-    """
-    input_text = f"<s> <<SYS>>{instruction}. {sys_message} <</SYS>> [INST]{prompt} [/INST]"
-    sequences = debias_pipeline(
-        input_text,
-        do_sample=True,
-        top_k=10,
-        num_return_sequences=1,
-        eos_token_id=debias_tokenizer.eos_token_id,
-        max_length=len(prompt.split(" ")) + len(input_text.split(" ")) + 100,  # Increased max_length
-    )
-    res = sequences[0]['generated_text']
-    result_part = res.split('[/INST]')[-1]
-    clean_result = ''.join(c for c in result_part if c.isprintable())
-    cleaned_text = re_incomplete_sentence(clean_result.strip())
-    return cleaned_text.strip()
-
+# Initialize the BiasPipeline
+bias_pipeline = BiasPipeline(model_loader)
 
 # Streamlit interface setup
 st.title('UnBIAS App')
-# Initialize the BiasPipeline
-bias_pipeline = BiasPipeline()
-
-# List of predefined test sentences
-test_sentences = [
-    "Women are just too emotional to be leaders.",
-    "All young people are lazy and addicted to their phones.",
-    "People from that country are always dishonest and corrupt.",
-    "Men are inherently better at science and math than women.",
-    "The elderly are just burdens to society; they contribute nothing.",
-    "All poor people are lazy and don't want to work.",
-    "People who follow that religion are all terrorists.",
-    "Immigrants are taking all our jobs and ruining the country.",
-    "All wealthy people achieved success only through manipulation and deceit.",
-    "People with tattoos are not professional.",
-    "All politicians are liars and cannot be trusted.",
-    "Blue-collar workers don't have any real skills.",
-    "Anyone who doesn't attend college is unintelligent.",
-    "Only people from the city are sophisticated.",
-    "Rural folks have backward mindsets and don't understand progress.",
-    "All artists are always broke and unreliable.",
-    "Anyone from that region is uneducated and narrow-minded.",
-    "People without jobs are simply not trying hard enough."
-]
 
 # Dropdown to select a test sentence
 selected_sentence = st.selectbox("Choose a pre-loaded sentence to analyze and debias:", [""] + test_sentences)