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()

# Streamlit interface
st.title('Text Debiasing App')
input_text = st.text_area("Enter text to debias:", height=150)
if st.button("Debias Text"):
    if input_text:
        debiased_text = get_debiased_sequence(input_text)
        st.write("Debiased Text:", debiased_text)
    else:
        st.write("Please enter some text to debias.")