Upload rlhf_refinement.py

rlhf_refinement.py

@@ -0,0 +1,122 @@
+import pandas as pd
+import numpy as np
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+import warnings
+warnings.filterwarnings("ignore")
+
+# Load the human evaluation dataset
+df = pd.read_excel("final_comments_evaluations_latest.xlsx")
+
+# Initialize the Granite 3.2-2B-Instruct model and tokenizer (from your existing setup)
+model_name = "ibm-granite/granite-3.2-2b-instruct"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
+
+# Define a simple reward model (mockup based on dataset metrics)
+# In practice, this would be the trained reward model from Stage 3
+def reward_model(paraphrase, original_scores):
+    # Mock reward calculation: adjust scores based on trends in the dataset
+    base_toxicity = original_scores["toxicity"]
+    base_empathy = original_scores["empathy"]
+    # Simulate improved paraphrasing: reduce toxicity, increase empathy
+    new_toxicity = max(0.1, base_toxicity - 0.2)  # Reduce toxicity
+    new_empathy = min(0.9, base_empathy + 0.1)   # Increase empathy
+    new_bias = original_scores["bias"]
+    new_hallucination = max(0.1, original_scores["hallucination"] - 0.1)
+    # Composite reward score (weights based on dataset analysis)
+    reward = 0.4 * new_empathy - 0.3 * new_toxicity - 0.2 * new_bias - 0.1 * new_hallucination
+    return reward, {"toxicity": new_toxicity, "empathy": new_empathy, "bias": new_bias, "hallucination": new_hallucination}
+
+# Function to generate a paraphrase using your existing paraphrasing logic
+def generate_paraphrase(comment, max_length=128):
+    prompt = (
+        "You are a content moderator tasked with rewriting toxic comments into neutral and constructive ones while maintaining the original meaning. "
+        "Follow these guidelines:\n"
+        "- Remove explicit hate speech, personal attacks, or offensive language.\n"
+        "- Keep the response neutral and professional.\n"
+        "- Ensure the rewritten comment retains the original intent but in a constructive tone.\n"
+        "- Match the length and brevity of the original toxic comment whenever possible. Keep the response short and to the point.\n\n"
+        "Examples:\n"
+        "Toxic: \"You're so dumb! You never understand anything!\"\n"
+        "Neutral: \"You might be misunderstanding this.\"\n"
+        "Toxic: \"This is the worst idea ever. Only an idiot would suggest this.\"\n"
+        "Neutral: \"I don’t think this idea works well.\"\n"
+        "Toxic: \"You’re useless.\"\n"
+        "Neutral: \"This isn’t helping much.\"\n"
+        "Toxic: \"Shut up.\"\n"
+        "Neutral: \"Let’s take a break from this.\"\n\n"
+        f"Now, rewrite this comment: \"{comment}\""
+    )
+    inputs = tokenizer(prompt, return_tensors="pt", max_length=max_length, truncation=True).to(device)
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=50,
+        num_beams=4,
+        early_stopping=True,
+        do_sample=False
+    )
+    paraphrase = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # Clean up the output by removing the prompt part
+    paraphrase = paraphrase.replace(prompt, "").strip()
+    if paraphrase.startswith("Neutral: "):
+        paraphrase = paraphrase[len("Neutral: "):].strip()
+    return paraphrase
+
+# RLHF Loop
+max_iterations = 5
+reward_threshold = 0.2  # Target for acceptable paraphrases (based on dataset range -0.25 to 0.24)
+results = []
+
+for idx, row in df.iterrows():
+    original_comment = row["Comment"]
+    current_paraphrase = row["Paraphrase_Comment"]
+    current_reward = row["reward_score"]
+    current_scores = {
+        "toxicity": row["toxicity"],
+        "empathy": row["empathy"],
+        "bias": row["bias"],
+        "hallucination": row["hallucination"]
+    }
+    
+    best_paraphrase = current_paraphrase
+    best_reward = current_reward
+    best_scores = current_scores.copy()
+    
+    # Iteratively refine the paraphrase
+    for iteration in range(max_iterations):
+        # Generate a new paraphrase
+        new_paraphrase = generate_paraphrase(original_comment)
+        # Evaluate the new paraphrase with the reward model
+        new_reward, new_scores = reward_model(new_paraphrase, current_scores)
+        
+        # If the new reward is better, update the best paraphrase
+        if new_reward > best_reward:
+            best_paraphrase = new_paraphrase
+            best_reward = new_reward
+            best_scores = new_scores
+        
+        # Stop if the reward exceeds the threshold
+        if best_reward >= reward_threshold:
+            break
+    
+    # Store the result
+    results.append({
+        "Comment": original_comment,
+        "Original_Paraphrase": current_paraphrase,
+        "Refined_Paraphrase": best_paraphrase,
+        "Original_Reward_Score": current_reward,
+        "Refined_Reward_Score": best_reward,
+        "Refined_Empathy": best_scores["empathy"],
+        "Refined_Toxicity": best_scores["toxicity"],
+        "Refined_Bias": best_scores["bias"],
+        "Refined_Hallucination": best_scores["hallucination"],
+        "Human_Evaluation_Reasoning": row["Human_Evaluation_Reasoning"]
+    })
+
+# Save the results to a CSV file
+results_df = pd.DataFrame(results)
+results_df.to_csv("refined_paraphrases.csv", index=False)
+print("Refinement complete. Results saved to refined_paraphrases.csv")