experiment.py

import pandas as pd
import phoenix as px
import nest_asyncio

from phoenix.experiments import run_experiment
from phoenix.evals import (
    llm_classify,
    LiteLLMModel,
    HALLUCINATION_PROMPT_RAILS_MAP,
    HALLUCINATION_PROMPT_TEMPLATE,
    QA_PROMPT_RAILS_MAP,
    QA_PROMPT_TEMPLATE,
)


from phoenix.trace import SpanEvaluations
from opentelemetry.trace import (
    format_span_id,
    get_current_span,
)


from phoenix.otel import register

PROJECT_NAME = "test"
register(
    project_name=PROJECT_NAME,
    auto_instrument=True,
)

model = LiteLLMModel(
    model="openrouter/openai/gpt-4o",
    temperature=0.0,
)

nest_asyncio.apply()


def task(expected):
    return expected["answer"]


def hallucination(input, expected) -> bool:
    data = pd.DataFrame(
        {
            "input": [input["question"]],
            "reference": [expected["reference"]],
            "output": [expected["answer"]],
        }
    )

    rails = list(HALLUCINATION_PROMPT_RAILS_MAP.values())
    eval_df = llm_classify(
        data=data,
        template=HALLUCINATION_PROMPT_TEMPLATE,
        model=model,
        rails=rails,
        provide_explanation=True,
    )

    eval_df["score"] = eval_df.apply(
        lambda x: 1 if x["label"] == "factual" else 0, axis=1
    )

    span = get_current_span()
    span_id = format_span_id(span.get_span_context().span_id)
    eval_df["span_id"] = span_id

    px.Client().log_evaluations(
        SpanEvaluations(
            dataframe=eval_df,
            eval_name="Hallucination",
        ),
    )

    return eval_df["label"].iloc[0] == "factual"


def correctness(input, expected) -> bool:
    data = pd.DataFrame(
        {
            "input": [input["question"]],
            "reference": [expected["reference"]],
            "output": [expected["answer"]],
        }
    )

    rails = list(QA_PROMPT_RAILS_MAP.values())
    eval_df = llm_classify(
        data=data,
        template=QA_PROMPT_TEMPLATE,
        model=model,
        rails=rails,
        provide_explanation=True,
    )

    eval_df["score"] = eval_df.apply(
        lambda x: 1 if x["label"] == "correct" else 0, axis=1
    )

    span = get_current_span()
    span_id = format_span_id(span.get_span_context().span_id)
    eval_df["span_id"] = span_id

    px.Client().log_evaluations(
        SpanEvaluations(
            dataframe=eval_df,
            eval_name="Q&A Correctness",
        ),
    )

    return eval_df["label"].iloc[0] == "correct"


dataset = px.Client().get_dataset(name="questions-v1")
run_experiment(
    dataset,
    task=task,
    evaluators=[hallucination, correctness],
    experiment_name="Overall experiment",
    experiment_description="Evaluating the overall experiment",
    experiment_metadata={"vendor": "openai", "model": "gpt-4o"},
)