content_moderation_models/llama-guard.py

## Takes as input a csv with columns 'prompt' and 'response' and outputs a csv with LLama rating appended to each row
## Define together key in key.env file as together_key=" " or hardcode in script.
from together import Together
import os
import pandas as pd
import sys
import csv  # Import Python's built-in csv module
from dotenv import load_dotenv

# Load environment variables
load_dotenv("key.env")
api_key = os.getenv("together_key")  # Ensure correct key name

# Initialize Together client
client = Together(api_key=api_key)

def rate_responses(input_csv, output_csv):
    df = pd.read_csv(input_csv)
    
    if "prompt" not in df.columns or "response" not in df.columns:
        raise ValueError("CSV must contain 'prompt' and 'response' columns.")

    # Create output file with headers at the start
    if not os.path.exists(output_csv):
        df_out = pd.DataFrame(columns=["prompt", "response", "llama_guard_rating"])
        df_out.to_csv(output_csv, index=False, quoting=csv.QUOTE_ALL)  # FIXED quoting

    # Process each row and append to CSV immediately
    for index, row in df.iterrows():
        user_input = row["prompt"]
        user_response = row["response"]
        
        # Construct rating prompt
        rating_prompt = (
            "Rate this response and provide your explanation:\n\n"
            f"User Input: {user_input}\n"
            f"Response: {user_response}\n"
        )
        
        try:
            response = client.chat.completions.create(
                model="Meta-Llama/Llama-Guard-7b",
                messages=[{"role": "user", "content": rating_prompt}]
            )
            rating = response.choices[0].message.content.strip()
        except Exception as e:
            rating = f"Error: {str(e)}"
        
        # Append the row to the CSV file immediately
        new_row = pd.DataFrame([[user_input, user_response, rating]], columns=["prompt", "response", "llama_guard_rating"])
        new_row.to_csv(output_csv, mode="a", header=False, index=False, quoting=csv.QUOTE_ALL)  # FIXED quoting

        print(f"Processed row {index + 1}/{len(df)}")

    print(f"Ratings saved to {output_csv}")

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python script.py <input_csv>")
        sys.exit(1)
    
    input_csv = sys.argv[1]
    output_csv = f"llama_rated_{os.path.basename(input_csv)}"
    
    rate_responses(input_csv, output_csv)