Polar Llama

Basic usage

POLAR_LLAMA_LOCAL_KV_BITS=4 python tag_column.py

import json
import dotenv
import polars as pl
from pydantic import BaseModel

from polar_llama import (
    Provider,
    combine_messages,
    execute_tool_calls,
    inference_messages,
    tool_results_to_message,
    tools_to_response_model,
)

dotenv.load_dotenv()

TOOLS = [
    {
        "name": "search_food_db",
        "description": "Search a nutrition database for one food item and portion.",
        "input_schema": {
            "type": "object",
            "properties": {
                "query": {"type": "string", "description": "The food item, singular"},
                "portion": {"type": "string", "description": "Portion as stated, e.g. 'two eggs'"},
            },
            "required": ["query", "portion"],
        },
    }
]

FAKE_DB = {
    "egg": 78, "sourdough toast": 120, "butter": 102, "black coffee": 2,
    "chicken bowl": 510, "guacamole": 230, "rice": 200,
}

def nutrition_executor(tool_name: str, arguments: dict):
    query = arguments["query"].lower()
    for food, calories in FAKE_DB.items():
        if food in query or query in food:
            return json.dumps({"food": food, "calories_per_serving": calories,
                               "portion": arguments["portion"]})
    return json.dumps({"food": query, "error": "not found"}), True

class NutritionSummary(BaseModel):
    total_calories: int
    items: list[str]
    notes: str

meals = pl.DataFrame({"meal": [
    "two eggs, sourdough toast with butter, black coffee",
    "chipotle chicken bowl, no rice, extra guac",
]})

FoodSearches = tools_to_response_model(TOOLS)

result = (
    meals
    # Turn 1a: the LLM parameterizes N searches per row (structured output).
    .with_columns(
        calls=pl.col("meal").llama.inference_async(
            provider=Provider.OPENAI, model="gpt-4o-mini",
            response_model=FoodSearches,
        )
    )
    # Turn 1b: all searches across all rows execute concurrently.
    .with_columns(
        results=execute_tool_calls(pl.col("calls"), executor=nutrition_executor, tools=TOOLS)
    )
    # Turn 2: reconcile results into a typed summary.
    .with_columns(
        nutrition=inference_messages(
            combine_messages(
                pl.col("meal").llama.to_message(role="user"),
                tool_results_to_message(pl.col("results")),
            ),
            provider=Provider.OPENAI, model="gpt-4o-mini",
            response_model=NutritionSummary,
        )
    )
)

print(result.select("meal", "nutrition"))

import polars as pl
from polar_llama import Predict, Signature, BootstrapFewShot, InstructionOptimizer, evaluate

# 1. Declare the task
module = Predict(
    Signature("question -> answer", instructions="Answer concisely."),
    provider="openai",
    model="gpt-4o-mini",
)

# 2. Labeled training data
trainset = pl.DataFrame({
    "question": ["What is 2+2?", "Capital of France?", "Largest planet?"],
    "answer": ["4", "Paris", "Jupiter"],
})

# 3. A metric: (gold row, prediction) -> bool | float
def exact_match(example, prediction):
    return example["answer"].strip().lower() == (prediction["answer"] or "").strip().lower()

# 4a. Bootstrap few-shot demos from rows the model already gets right
compiled = BootstrapFewShot(metric=exact_match, max_demos=4).compile(module, trainset)

# 4b. Or search for better instructions (COPRO-style)
optimizer = InstructionOptimizer(metric=exact_match, n_candidates=4)
compiled = optimizer.compile(module, trainset)
print(optimizer.history)  # [(instructions, score), ...]

# 5. Run the optimized module on new data — outputs land in pred_* columns
result = compiled(pl.DataFrame({"question": ["What is 3+3?"]}))
print(result["pred_answer"])

# Score any module against a labeled set
print(evaluate(compiled, trainset, exact_match).score)

import polars as pl
from polar_llama import combine_messages, inference_messages
import dotenv

dotenv.load_dotenv()

df = pl.DataFrame({
    "system_prompt": [
        "You are a helpful assistant.",
        "You are a math expert."
    ],
    "user_question": [
        "What's the weather like today?",
        "Solve x^2 + 5x + 6 = 0"
    ]
})

# Using .llama namespace (recommended)
df = df.with_columns([
    pl.col("system_prompt").llama.to_message(role="system").alias("system_message"),
    pl.col("user_question").llama.to_message(role="user").alias("user_message")
])

df = df.with_columns(
    conversation=combine_messages(pl.col("system_message"), pl.col("user_message"))
)

df = df.with_columns(
    response=inference_messages(pl.col("conversation"), provider="openai", model="gpt-4o")
)

print(df.select(["user_question", "response"]))

import polars as pl
from polar_llama import Provider

feedback_df = pl.DataFrame({
    'customer_id': [101, 102, 103, 104, 105],
    'feedback': [
        'The product is amazing but shipping was slow',
        'Great quality, highly recommend!',
        'Disappointed with customer service',
        'Perfect for my needs, will buy again',
        'Product arrived damaged, requesting refund'
    ]
})

sentiment_prompt = """Analyze the sentiment of this customer feedback
and classify it as Positive, Negative, or Neutral.
Also provide a brief reason.

Feedback: {feedback}"""

df = feedback_df.with_columns(
    prompt=pl.format(sentiment_prompt, pl.col('feedback'))
)

df = df.with_columns(
    sentiment_analysis=pl.col('prompt').llama.inference_async(
        provider=Provider.OPENAI,
        model='gpt-4o-mini'
    )
)

print(df.select(['customer_id', 'feedback', 'sentiment_analysis']))