AI guardrails: Relevance scorers

HelpSteer2 Relevance =(0.4×Helpfulness)+(0.3×Correctness)+(0.2×Coherence)−(0.1×Verbosity)

import asyncio
import weave; weave.init("relevance-scorer")

from weave.scorers import RelevanceScorer


# Initialize the relevance scorer
relevance_scorer = RelevanceScorer()

# Define input, context, and output
input_text = "Some input text to consider"
context_text = "Some context to consider"  # Optional context
output_text = "Some output text to consider"

# Async function to compute relevance score
async def get_relevance_score():
    score = await relevance_scorer.score(input=input_text, context=context_text, output=output_text)
    print(score)

# Run the async function
asyncio.run(get_relevance_score())

import weave
from weave.scorers import CoherenceScorer
import pandas as pd
from datasets import load_dataset
import asyncio
from weave.trace.box import unbox
import os
from pydantic import BaseModel, Field
from typing import Literal
from typing import Optional, Any
from litellm import acompletion
import json 
# from relevance_scorer import RelevanceScorer
from weave.scorers import RelevanceScorer
weave.init("relevance_eval")



def load_relevance_dataset():
    # Load the HelpSteer2 dataset
    dataset = load_dataset("nvidia/HelpSteer2", split="validation")
    
    # Convert to a Pandas DataFrame for easier manipulation
    df = pd.DataFrame(dataset)
    
    # Define the weights for relevance computation
    weights = {
        "helpfulness": 0.4,
        "correctness": 0.3,
        "coherence": 0.2,
        "verbosity": -0.1  # Negative weight for verbosity
    }

    # Calculate relevance score using the weighted formula
    df["relevance"] = (
        weights["helpfulness"] * df["helpfulness"] +
        weights["correctness"] * df["correctness"] +
        weights["coherence"] * df["coherence"] +
        weights["verbosity"] * df["verbosity"]
    )
    
    # Determine binary labels (True for relevance >= 3.5, False otherwise)
    df["relevance_label"] = df["relevance"] >= 3.5
    
    # Separate into True and False groups
    true_group = df[df["relevance_label"] == True]
    false_group = df[df["relevance_label"] == False]

    # Determine the maximum number of samples per group
    max_samples = min(len(true_group), len(false_group), 75)  # Adjust dynamically based on available data
    
    # Sample from each group
    balanced_true = true_group.sample(n=max_samples, random_state=42)
    balanced_false = false_group.sample(n=max_samples, random_state=42)
    
    # Combine the balanced samples
    balanced_df = pd.concat([balanced_true, balanced_false]).reset_index(drop=True)
    
    # Shuffle the resulting DataFrame
    balanced_df = balanced_df.sample(frac=1, random_state=42).reset_index(drop=True)
    
    # Prepare the dataset
    dataset_prepared = [
        {"output": row["response"], "label": row["relevance_label"], "input": row["prompt"]}
        for _, row in balanced_df.iterrows()
    ]
    
    return dataset_prepared


RELEVANCE_SYSTEM_PROMPT = """You are an expert evaluator assessing the relevance of LLM-generated outputs relative to their input context. 
Your goal is to provide a single relevance score and classification based on comprehensive analysis.
Relevance measures how effectively a generated output addresses its input context across three core dimensions:

1. **Semantic Alignment**
   - How directly does the output address key input requirements?
   - Does it maintain topical focus?
   - Does it provide complete coverage of necessary information?
   - Is unnecessary content avoided?

2. **Structural Coherence**
   - Does the output flow logically and show internal consistency?
   - Is the presentation of information clear and organized?
   - Is there a good balance between completeness and conciseness?

3. **Contextual Integration**
   - How well does the output use the provided context?
   - Does the output align with the broader discourse?
   - Is it consistent with background information?
   - Does it fulfill task-specific requirements?

## Evaluation Process

1. Review all input context (instructions, prompts, documents, chat history)
2. Identify core requirements and purpose
3. Analyze the LLM output across all three dimensions
4. Assign a single relevance score (1-5):
   - 5: Exceptional relevance across all dimensions
   - 4: Strong relevance with minor gaps
   - 3: Adequate relevance with some issues
   - 2: Significant relevance issues
   - 1: Major relevance problems
5. Classify as relevant (score ≥ 3.0) or not relevant (score < 3.0)

## Task-Specific Considerations

- **Summarization**: Focus on key information selection and density
- **Q&A**: Emphasize answer accuracy and completeness
- **Chat**: Consider conversation flow and context maintenance
- **RAG**: Evaluate retrieved information integration

## Output Format

Provide evaluation results in the following JSON format:

```json
{
  "relevance": [score from 1-5],
  "relevant": [true/false]
}
```
""".strip()

class RelevanceScore(BaseModel):
    """The level of relevance of a <completion> for a given <context>."""
    chain_of_thought: str = Field(..., description="The chain of thought that led to the prediction")
    score: int = Field(..., description="Score the relevance of the <completion> on a likert scale of 1 to 5")
    relevance: Literal["Perfectly Relevant", "Mostly Relevant", "A Little Irrelevant", "Mostly Irrelevant", "Completely Irrelevant"] = Field(..., description="The level of relevance of the <completion>")
    relevant: bool = Field(..., description="Whether the <completion> is relevant or not, anything above 3 is relevant")




class OpenAIRelevanceScorer(weave.Model):
    system_prompt: str = RELEVANCE_SYSTEM_PROMPT
    model_id: str 
    temperature: float = 0.7
    max_tokens: int = 4096
    # client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
    
    def __init__(self, model_id: str, temperature: float = 0.7, max_tokens: int = 4096):
        """
        Initialize the OpenAIRelevanceScorer with a specific model ID.
        """
        super().__init__(model_id=model_id, temperature=temperature, max_tokens=max_tokens)
        self.model_id = model_id
        self.temperature = temperature
        self.max_tokens = max_tokens
    def _format_messages(self, prompt: str, completion: str, context:Optional[list[str]], chat_history: Optional[list[dict[str, str]]]) -> str:
        """Format the prompt for the model."""


        chat_history = chat_history if isinstance(chat_history, list) else []
        context = context if isinstance(context, list) else []
        if context:
            context = "\n".join(context).strip()
            context = f"<documents>\n{context}\n</documents>"
        else:
            context = ""
        prompt = f"{context}\n\n{prompt}".strip()


        messages = chat_history + [{"role": "user", "content": prompt}]


        messages = [f"<|msg_start|>{message['role']}\n{message['content']}<|msg_end|>" for message in messages]
        messages = "\n".join(messages)
        
        context = f"<context>{messages}</context>\n"
        completion = f"<completion>{completion}</completion>\n"
        
        context_and_completion = context + completion
        
        return [
            {"role": "system", "content": self.system_prompt},
            {"role": "user", "content": context_and_completion},
            ]
    

    async def run_inference(self, messages):
        api_key = os.getenv("OPENAI_API_KEY")  # Ensure your API key is set in the environment variable
        model_name = self.model_id

        response = await acompletion(
            model=model_name,
            api_key=api_key,
            messages=messages,
            temperature=self.temperature,
            max_tokens=self.max_tokens,
        )
        
        return response["choices"][0]["message"]["content"]
    
    @weave.op
    async def predict(self, input: str, output:str, context:Optional[list[str]] = None, chat_history: Optional[list[dict[str, str]]] = None) -> dict[str, Any]:
        try: 
            messages = self._format_messages(prompt=input, completion=output, context=context, chat_history=chat_history)

            response = await self.run_inference(messages)
            if response.startswith('```') and response.endswith('```'):
                response = response.split('\n', 1)[1].rsplit('\n', 1)[0]        
            parsed_response = json.loads(response)  # Convert the string to a Python dictionary
            return bool({k: v for k, v in parsed_response.items() if k != "chain_of_thought"}['relevant'])
        except Exception as e:
            print(str(e))
            return False



# Define Weave Relevance Scorer Model
class WeaveRelevanceModel(weave.Model):
    @weave.op
    async def predict(self, input: str) -> int:
        result = await scorer.score(input=input, output=input)
        # return result['relevant'] 
        return result['score'] >= 3 # multiple ways to determine relevance prediction here 


class PrecisionRecallF1Scorer(weave.Scorer):
    """
    Custom scorer to calculate precision, recall, F1, and accuracy at the dataset level.
    """

    @weave.op
    def score(self, label: int, model_output: int) -> dict:
        """
        Compute True Positives, False Positives, False Negatives, and True Negatives for a single row.
        """
        tp = int(label == 1 and model_output == 1)  # True Positive
        fp = int(label == 0 and model_output == 1)  # False Positive
        fn = int(label == 1 and model_output == 0)  # False Negative
        tn = int(label == 0 and model_output == 0)  # True Negative

        return {"tp": tp, "fp": fp, "fn": fn, "tn": tn}

    def summarize(self, score_rows: list) -> dict:
        """
        Summarize precision, recall, F1, and accuracy from the row-level scores.
        """
        # Aggregate true positives, false positives, false negatives, and true negatives
        total_tp = sum(row["tp"] for row in score_rows)
        total_fp = sum(row["fp"] for row in score_rows)
        total_fn = sum(row["fn"] for row in score_rows)
        total_tn = sum(row["tn"] for row in score_rows)

        # Calculate precision, recall, F1, and accuracy
        precision = total_tp / (total_tp + total_fp) if (total_tp + total_fp) > 0 else 0
        recall = total_tp / (total_tp + total_fn) if (total_tp + total_fn) > 0 else 0
        f1 = (
            2 * (precision * recall) / (precision + recall)
            if (precision + recall) > 0
            else 0
        )
        accuracy = (total_tp + total_tn) / (total_tp + total_fp + total_fn + total_tn) if (total_tp + total_fp + total_fn + total_tn) > 0 else 0

        return {
            "precision": precision,
            "recall": recall,
            "f1": f1,
            "accuracy": accuracy,
            # "tp": total_tp,
            # "fp": total_fp,
            # "fn": total_fn,
            # "tn": total_tn,
        }



# Instantiate OpenAI Relevance Scorers
openai_scorer_4o = OpenAIRelevanceScorer(model_id="gpt-4o-2024-08-06", temperature=0.0, max_tokens=4096)
openai_scorer_4omini = OpenAIRelevanceScorer(model_id="gpt-4o-mini", temperature=0.0, max_tokens=2048)
scorer = RelevanceScorer()



# Define OpenAIScorer4o Model
class OpenAIScorer4o(weave.Model):
    """
    Weave model that wraps the OpenAIRelevanceScorer for gpt-4o.
    """
    @weave.op
    async def predict(self, input: str, output: str, context: Optional[list[str]] = None, chat_history: Optional[list[dict[str, str]]] = None) -> bool:
        """
        Use the instantiated gpt-4o scorer to predict relevance.
        """
        return await openai_scorer_4o.predict(input=input, output=output, context=context, chat_history=chat_history)


# Define OpenAIScorer4omini Model
class OpenAIScorer4omini(weave.Model):
    """
    Weave model that wraps the OpenAIRelevanceScorer for gpt-4o-mini.
    """
    @weave.op
    async def predict(self, input: str, output: str, context: Optional[list[str]] = None, chat_history: Optional[list[dict[str, str]]] = None) -> bool:
        """
        Use the instantiated gpt-4o-mini scorer to predict relevance.
        """
        return await openai_scorer_4omini.predict(input=input, output=output, context=context, chat_history=chat_history)

# Run the evaluations
async def run_evaluations():
    """Run evaluations for the  scorers."""
    # Load dataset
    dataset = load_relevance_dataset()
    print("Dataset loaded...")
    # Initialize models
    models = {
        "WeaveRelevanceModel": WeaveRelevanceModel(),
        "OpenAIRelevanceScorer4oMini": OpenAIScorer4omini(),  
        "OpenAIRelevanceScorer4o": OpenAIScorer4o(),  
    }

    # Define evaluation scorers
    scorers = [PrecisionRecallF1Scorer()]

    # Run evaluations
    results = {}
    for model_name, model in models.items():
        print(f"\nEvaluating {model_name}...")
        evaluation = weave.Evaluation(
            dataset=dataset,
            scorers=scorers,
            name=model_name + " Eval"
        )
        results[model_name] = await evaluation.evaluate(model)

    # Print results
    for model_name, result in results.items():
        print(f"\nResults for {model_name}:")
        print(result)


if __name__ == "__main__":
    asyncio.run(run_evaluations())