Update app.py

app.py

@@ -1,47 +1,57 @@
+%%writefile debias_app.py
 import streamlit as st
-from transformers import AutoTokenizer, pipeline
-import transformers
-import torch
-import pandas as pd
-
-# Model setup
-model = "newsmediabias/UnBIAS-LLama2-Debiaser-Chat-QLoRA"
-tokenizer = AutoTokenizer.from_pretrained(model)
-debias_pipeline = transformers.pipeline(
-    "text-generation",
-    model=model,
-    torch_dtype=torch.float16,
-    device_map="auto",
-)
-
-# Sample Instruction
-instruction = ("Instruction: As a helpful, respectful and trustworthy debiasing assistant, your "
-               "task is to receive a text and return its unbiased version, without adding any unrelated content "
-               "or additional outputs.")
-
-def get_debiased_sequence(prompt):
-    """Generate a debiased version of the provided text using the debiasing pipeline."""
-    input_text = f"<s> <<SYS>> {instruction} <</SYS>> [INST]{prompt} [/INST]"
-    sequences = debias_pipeline(
-        input_text,
-        do_sample=True,
-        top_k=10,
-        num_return_sequences=1,
-        eos_token_id=tokenizer.eos_token_id,
-        max_length=len(prompt)+100,
-    )
-    res = sequences[0]['generated_text']
-    result_part = res.split('[/INST]')[-1]
-    clean_result = ''.join(c for c in result_part if c.isprintable())
-    return clean_result.strip()
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModelForTokenClassification, pipeline
+
+# Define the BiasPipeline class with text processing methods
+class BiasPipeline:
+    def __init__(self):
+        # Load models and tokenizers
+        self.load_resources()
+
+    def load_resources(self):
+        """Load models and tokenizers."""
+        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-Named-Entity-Recognition")
+        self.ner_model = AutoModelForTokenClassification.from_pretrained("newsmediabias/UnBIAS-Named-Entity-Recognition")
+
+        self.classifier = pipeline("text-classification", model=self.classifier_model, tokenizer=self.classifier_tokenizer)
+        self.ner = pipeline("ner", model=self.ner_model, tokenizer=self.ner_tokenizer)
+
+    def clean_text(self, text):
+        """Clean up the text by removing any redundant spaces."""
+        return ' '.join(text.split())
+
+    def complete_sentence(self, text):
+        """If the text ends mid-sentence, remove all words after the last full stop."""
+        sentences = text.split(". ")
+        if len(sentences) > 1 and not sentences[-1].endswith("."):
+            return ". ".join(sentences[:-1]) + "."
+        return text
+
+    def create_token_limit(self, text):
+        words = text.split()
+        max_length = round(len(words) + 1.5 * len(words))
+        return max_length
+
+    def process(self, texts):
+        """Process texts to classify and find named entities."""
+        classification_results = self.classifier(texts)
+        ner_results = self.ner(texts)
+        return classification_results, ner_results
+
+# Initialize the BiasPipeline
+pipeline = BiasPipeline()
 
 # Streamlit interface
 st.title('UnBIAS App')
-input_text = st.text_area("Enter text to debias:", height=150)
-if st.button("Debias Text"):
+input_text = st.text_area("Enter text:", height=150)
+if st.button("Process Text"):
     if input_text:
-        debiased_text = get_debiased_sequence(input_text)
-        st.write("Debiased Text:", debiased_text)
+        cleaned_text = pipeline.clean_text(input_text)
+        classification_results, ner_results = pipeline.process(cleaned_text)
+        st.write("Classification Results:", classification_results)
+        st.write("Named Entity Recognition Results:", ner_results)
     else:
-        st.write("Please enter some text to debias.")
-
+        st.write("Please enter some text to process.")