import torch

class MyToxicityDebiaserPipeline(object):
    def __init__(self, model, tokenizer, gpt_model, gpt_tokenizer, device=None, **kwargs):
        self.model = model.to(device)
        self.tokenizer = tokenizer
        self.gpt_model = gpt_model.to(device)
        self.gpt_tokenizer = gpt_tokenizer
        self.device = device if device is not None else torch.device("cpu")

    def _forward(self, inputs):
        text = inputs["text"]
        encoded = self.tokenizer(text, truncation=True, padding=True, return_tensors="pt").to(self.device)
        logits = self.model(encoded.input_ids, encoded.attention_mask).logits
        probs = torch.softmax(logits, dim=-1)
        label = torch.argmax(probs, dim=-1).item()
        return {"label": label, "probabilities": probs.tolist(), "text_input_ids": encoded.input_ids}

    def _sanitize_parameters(self, **kwargs):
        return kwargs, {}, {}

    def preprocess(self, inputs):
        return {"text": inputs}

    def postprocess(self, outputs):
        label = outputs["label"]
        if label == 0:
            prompt = "This comment is non-toxic."
        elif label == 1:
            prompt = "This comment is toxic but has been debiased as follows:"
            text = self.tokenizer.decode(outputs["text_input_ids"][0])

            debias_prompt = f"Remove the offensive words and biased tone and write the same sentence nicely: {text}"
            encoded_debias_prompt = self.gpt_tokenizer.encode_plus(debias_prompt, return_tensors="pt").to(self.device)

            generated = self.gpt_model.generate(
                input_ids=encoded_debias_prompt["input_ids"],
                attention_mask=encoded_debias_prompt["attention_mask"],
                do_sample=True,
                max_length=100,
                top_p=0.95,
                temperature=0.7,
                pad_token_id=self.gpt_tokenizer.pad_token_id,
                eos_token_id=self.gpt_tokenizer.eos_token_id,
                early_stopping=True,
            )
            generated_text = self.gpt_tokenizer.decode(generated[0], skip_special_tokens=True)
            prompt += f"\nOriginal text: {text}\nDebiased text: {generated_text}"
        return prompt


    def __call__(self, inputs, *args, **kwargs):
        _args, _kwargs, forward_kwargs = self._sanitize_parameters(*args, **kwargs)
        inputs = self.preprocess(inputs)
        outputs = self._forward(inputs, **forward_kwargs)
        return self.postprocess(outputs)