Running inference and evaluating Llama 4 in Python

import os
import torch
from openai import OpenAI
from transformers import AutoTokenizer, Llama4ForConditionalGeneration, FbgemmFp8Config
import weave; weave.init("llama4")

OPENROUTER_API_KEY = ""
USE_OPENROUTER = True if OPENROUTER_API_KEY else False

def infer_with_openrouter(prompt):
    client = OpenAI(
        base_url="https://openrouter.ai/api/v1",
        api_key=OPENROUTER_API_KEY
    )
    messages = [{"role": "user", "content": prompt}]
    response = client.chat.completions.create(
        model="meta-llama/llama-4-scout",
        messages=messages,
        max_tokens=300
    )
    return response.choices[0].message.content

def infer_with_huggingface(prompt):
    model_id = "meta-llama/Llama-4-Scout-17B-16E-Instruct"
    tokenizer = AutoTokenizer.from_pretrained(model_id)

    if not torch.cuda.is_available():
        print("WARNING: CUDA not available. This will be slow as hell or might just crash.")

    model = Llama4ForConditionalGeneration.from_pretrained(
        model_id,
        torch_dtype=torch.bfloat16,
        quantization_config=FbgemmFp8Config(),
        device_map="auto" if torch.cuda.is_available() else None
    )

    messages = [{"role": "user", "content": prompt}]
    inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, max_new_tokens=300)
    return tokenizer.decode(outputs[0, inputs.shape[-1]:], skip_special_tokens=True)
@weave.op 
def infer(prompt):
    if USE_OPENROUTER:
        print("[ using OpenRouter ]")
        return infer_with_openrouter(prompt)
    else:
        print("[ using HuggingFace ]")
        return infer_with_huggingface(prompt)

# example
if __name__ == "__main__":
    output = infer("What is the capital of France?")
    print("\n" + output)


import base64
import random
import asyncio
import os
from datasets import load_dataset
from io import BytesIO
from PIL import Image
from openai import OpenAI
import weave
from weave import Evaluation, Model
import json 
from litellm import completion
import time

weave.init("meta_llama_eval")

OPENAI_API_KEY = "yourapikey"  # Replace with your actual OpenAI API key
meta_client = OpenAI(
    base_url="https://openrouter.ai/api/v1",
    api_key="your_openrouter_apikey",
)


SEED = 3
random.seed(SEED)

# Function to perform inference using litellm for the scorer
def run_inference_openai(prompt, model_id="gpt-4o-2024-08-06"):
    try:
        response = completion(
            model=model_id,
            temperature=0.0,
            messages=[
                {"role": "user", "content": prompt}
            ]
        )
        
        # Extract content from litellm response
        if response and hasattr(response, 'choices') and len(response.choices) > 0:
            content = response.choices[0].message.content
            return content
        else:
            print("No content found in response")
            return None
            
    except Exception as e:
        print(f"Failed to get response: {e}")
        return None

# Function to encode PIL image to base64
def encode_pil_image_to_base64(pil_image):
    buffered = BytesIO()
    pil_image.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode("utf-8")

@weave.op
async def gpt4o_scorer(expected: str, model_output: str) -> dict:
    """Score the model's output by comparing it with the ground truth."""
    query = f"""
    YOU ARE A LLM JUDGE DETERMINING IF THE FOLLOWING MODEL GENERATED ANSWER IS THE SAME AS THE CORRECT ANSWER FOR THE CHART-QA DATASET.
    It's ok if predicted/ground truth answers aren't formatted the same or slightly wored differently 

    Model's Answer: {str(model_output)}
    Correct Answer: {expected}
    
    Your task:
    1. State the model's predicted answer (answer only).
    2. State the ground truth (answer only).
    3. Determine if the model's final answer is correct (ignore formatting differences, etc.). RESPOND with the predicted and ground truth answer, followed with a JSON object containing the correctness encapsulated within the following delimiters:
       ```json
       {{ "correctness": true/false }}
       ```
    """
    
    # Perform inference using litellm
    response = run_inference_openai(query, "gpt-4o-2024-08-06")
    
    if response is None:
        return {"correctness": False, "reasoning": "Inference failed."}
    
    try:
        # Extract correctness JSON object from the response
        json_start = response.index("```json") + 7
        json_end = response.index("```", json_start)
        correctness = json.loads(response[json_start:json_end].strip()).get("correctness", False)
    except (ValueError, IndexError):
        correctness = False

    return {"correctness": correctness, "reasoning": response}

# Load dataset and preprocess
def load_ds():
    dataset = load_dataset("TeeA/ChartQA", split="test", cache_dir="./cache")
    dataset = dataset.shuffle(seed=SEED)
    eval_data = dataset.select(range(len(dataset) - 100, len(dataset)))

    def sample_image_and_query(example):
        image = example["image"]
        img = Image.open(image) if isinstance(image, str) else image
        return {
            "question": example["qa"][0]["query"],
            "image": img,  # Pass the PIL image directly
            "expected": example["qa"][0]["label"]
        }

    return [sample_image_and_query(example) for example in eval_data]

# Scoring function
@weave.op()
def substring_match(expected: str, model_output: dict) -> dict:
    match = expected.lower() in model_output['output'].lower()
    return {"substring_match": match}

# Model classes - now converting PIL images to base64 inside the predict methods
class MetaLLaMA4Scout(Model):
    @weave.op()
    def predict(self, question: str, image: Image.Image):
        # Convert PIL image to base64 inside the predict method
        image_base64 = f"data:image/png;base64,{encode_pil_image_to_base64(image)}"
        
        time.sleep(2)  # delay in seconds

        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": question},
                    {"type": "image_url", "image_url": {"url": image_base64}}
                ]
            }
        ]
        response = meta_client.chat.completions.create(
            model="meta-llama/llama-4-scout",
            messages=messages,
            max_tokens=300,
            temperature=0.0
        )
        return {"output": response.choices[0].message.content}

class MetaLLaMA4Maverick(Model):
    @weave.op()
    def predict(self, question: str, image: Image.Image):
        # Convert PIL image to base64 inside the predict method
        image_base64 = f"data:image/png;base64,{encode_pil_image_to_base64(image)}"
        
        time.sleep(2)  # delay in seconds

        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": question},
                    {"type": "image_url", "image_url": {"url": image_base64}}
                ]
            }
        ]
        response = meta_client.chat.completions.create(
            model="meta-llama/llama-4-maverick:free",
            messages=messages,
            max_tokens=300,
            temperature=0.0
        )
        return {"output": response.choices[0].message.content}

class GPT4oModel(Model):
    @weave.op()
    def predict(self, question: str, image: Image.Image):
        # Convert PIL image to base64 inside the predict method
        image_base64 = f"data:image/png;base64,{encode_pil_image_to_base64(image)}"
        
        time.sleep(2)
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": question},
                    {"type": "image_url", "image_url": {"url": image_base64}}
                ]
            }
        ]
        try:
            response = completion(
                model="openai/gpt-4o-2024-08-06",
                messages=messages,
                temperature=0
            )
            return {"output": response.choices[0].message.content}
        except Exception as e:
            return {"output": f"Error: {e}"}

# Main eval loop
async def run_evaluations():
    print("Loading dataset...")
    dataset = load_ds()
    print(f"Loaded {len(dataset)} problems")

    print("Initializing models...")
    models = {
        "llama-4-scout": MetaLLaMA4Scout(),
        "llama-4-maverick": MetaLLaMA4Maverick(),
        "gpt-4o": GPT4oModel()
    }

    print("Preparing dataset for evaluation...")
    dataset_prepared = dataset

    print("Running evaluations...")
    scorers = [gpt4o_scorer]

    for model_name, model in models.items():
        print(f"\n\n=== EVALUATING {model_name.upper()} ===")
        evaluation = Evaluation(
            dataset=dataset_prepared,
            scorers=scorers,
            name=f"{model_name} Evaluation"
        )
        results = await evaluation.evaluate(model)
        print(f"Results for {model_name}: {results}")

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