Training Search Agents

-- fetch_patents.sql
SELECT
  publication_number,
  title,
  abstract,
  claims,
  description,
  filing_date,
  grant_date,
  ARRAY_LENGTH(claims_localized) AS claim_count
FROM
  `patents-public-data.patents.publications`
WHERE
  assignee_harmonized.name IN (
    'QUALCOMM INCORPORATED',
    'ERICSSON',
    'NOKIA', 
    'SAMSUNG',
    'HUAWEI',
    'SONY',
  )
  AND country_code = 'US'
  AND grant_date IS NOT NULL
LIMIT 1500

# format_patents.py
def format_patent(row):
    doc = f"""# {row['title']}

## Abstract
{row['abstract']}

## Claims
{row['claims']}

## Description
{row['description']}
"""
    metadata = {
        "patent_id": row["publication_number"],
        "title": row["title"],
        "filing_date": row["filing_date"],
        "grant_date": row["grant_date"],
        "claim_count": row["claim_count"],
    }
    return doc, metadata

# generate_qa.py
import random
from datetime import datetime

def generate_date_question(patents: list[dict]) -> dict:
    patent = random.choice(patents)
    return {
        "question": f"When was patent {patent['patent_id']} filed?",
        "answer": patent["filing_date"],  # e.g. "2019-03-15"
        "patent_ids": [patent["patent_id"]],
        "level": 1,
        "type": "filing_date",
    }

def generate_days_to_grant(patents: list[dict]) -> dict:
    patent = random.choice([p for p in patents if p["grant_date"]])
    filed = datetime.strptime(patent["filing_date"], "%Y-%m-%d")
    granted = datetime.strptime(patent["grant_date"], "%Y-%m-%d")
    days = (granted - filed).days
    return {
        "question": f"How many days elapsed between the filing and grant of patent {patent['patent_id']}?",
        "answer": str(days),
        "patent_ids": [patent["patent_id"]],
        "level": 2,
        "type": "days_to_grant",
    }

# generate_l3_qa.py
import json
from openai import OpenAI

client = OpenAI()

GENERATION_PROMPT = """You are generating ground truth for a patent analysis training dataset.

Patent text:
{patent_text}

Question: {question}

Return ONLY valid JSON with no preamble:
{{
    "answer": "prose answer to the question",
    "key_points": ["specific factual claim 1", "specific factual claim 2"],
    "source_quotes": ["exact quote from patent supporting point 1", "exact quote supporting point 2"]
}}"""

def generate_l3_ground_truth(patent_text: str, question: str) -> dict:
    response = client.chat.completions.create(
        model="model",
        messages=[{
            "role": "user",
            "content": GENERATION_PROMPT.format(
                patent_text=patent_text,
                question=question
            )
        }],
        temperature=0
    )
    return json.loads(response.choices[0].message.content)

# rewards.py
def reward_exact(agent_answer: str, ground_truth: str) -> float:
    """Binary reward for single-answer questions."""
    return 1.0 if normalize(agent_answer) == normalize(ground_truth) else 0.0

def reward_set(agent_answer: str, valid_set: set[str]) -> float:
    """For search questions with multiple valid answers."""
    return 1.0 if normalize(agent_answer) in valid_set else 0.0

def normalize(s: str) -> str:
    return s.strip().lower().replace(",", "").replace(".", "")

// judge_output_schema.json
{
    "key_points": [
        {"point": "SSB spatial overloading concept", "covered": true},
        {"point": "Spatially separated beams", "covered": true},
        {"point": "Reduced bandwidth overhead", "covered": false}
    ],
    "hallucination": false,
    "factual_error": false,
    "final_score": 7
}

# rewards.py
def reward_l3(judge_output: dict) -> float:
    score = judge_output["final_score"] / 10.0  # normalize to [0, 1]
    if judge_output.get("hallucination"):
        score *= 0.2  # hallucination is catastrophic in patent context
    elif judge_output.get("factual_error"):
        score *= 0.5
    return score

# patent_env.py
import chromadb
from verifiers import ToolEnv

class PatentEnv(ToolEnv):
    def __init__(self, patents: list[dict], level: int):
        self.corpus = {p["patent_id"]: p for p in patents}
        self.level = level
        self._init_vector_store(patents)
        super().__init__(tools=self._get_tools(level))

    def _init_vector_store(self, patents):
        self.chroma = chromadb.Client()
        collection = self.chroma.create_collection("patents")
        collection.add(
            ids=[p["patent_id"] for p in patents],
            documents=[f"{p['title']} {p['abstract']}" for p in patents],
        )
        self.collection = collection

    def search_patents(self, query: str) -> list[dict]:
        """Returns patent IDs and titles matching a semantic query."""
        results = self.collection.query(query_texts=[query], n_results=5)
        return [
            {"patent_id": id_, "title": self.corpus[id_]["title"]}
            for id_ in results["ids"][0]
        ]

    def get_metadata(self, patent_id: str) -> dict:
        """Returns title, filing_date, grant_date, claim_count."""
        p = self.corpus[patent_id]
        return {
            "title": p["title"],
            "filing_date": p["filing_date"],
            "grant_date": p["grant_date"],
            "claim_count": p["claim_count"],
        }

    def get_abstract(self, patent_id: str) -> str:
        return self.corpus[patent_id]["abstract"]

    def view_sections(self, patent_id: str) -> list[str]:
        """Lists available section names in the full patent text."""
        doc = self.corpus[patent_id]["full_text"]
        return [
            line.lstrip("# ").strip()
            for line in doc.split("\n")
            if line.startswith("## ")
        ]

    def read_section(self, patent_id: str, section_name: str) -> str:
        """Returns the full text of a named section."""
        doc = self.corpus[patent_id]["full_text"]
        start = doc.find(f"## {section_name}")
        if start == -1:
            return f"Section '{section_name}' not found."
        end = doc.find("\n## ", start + 1)
        return doc[start:end if end != -1 else len(doc)]

    def _get_tools(self, level: int):
        tools = [self.search_patents, self.get_metadata]
        if level >= 2:
            tools.append(self.get_abstract)
        if level >= 3:
            tools += [self.view_sections, self.read_section]
        return tools

# parsing.py
import json

def extract_claims_text(full_text: str) -> str:
    """Extract the Claims section using string slicing, not regex."""
    start = full_text.find("## Claims")
    if start == -1:
        return ""
    end = full_text.find("\n## ", start + 1)
    return full_text[start:end if end != -1 else len(full_text)]

def parse_patent_id(full_code: str) -> str:
    """'US10123456B2' -> split and handle deterministically."""
    return full_code.split("-")[0] if "-" in full_code else full_code

def parse_judge_output(raw: str) -> dict:
    """Parse judge JSON with a robust fallback."""
    try:
        return json.loads(raw)
    except json.JSONDecodeError:
        # Fallback: find final_score by character iteration
        key = '"final_score":'
        idx = raw.find(key)
        if idx == -1:
            return {"final_score": 0, "hallucination": True, "key_points": []}
        score_start = idx + len(key)
        digits = ""
        for ch in raw[score_start:].lstrip():
            if ch.isdigit():
                digits += ch
            elif digits:
                break
        return {
            "final_score": int(digits) if digits else 0,
            "hallucination": "hallucin" in raw.lower(),
            "key_points": [],
        }

# Push each level as a separate named environment
prime env push --name basic-patent-q-and-a
prime env push --name advanced-patent-q-and-a
prime env push --name patent-technical-analysis

# config.toml
model = "meta-llama/Llama-3.2-3B-Instruct"
max_steps = 100
batch_size = 128
rollouts_per_example = 8

[sampling]
max_tokens = 1024

[[env]]
id = "primeintellect/advanced-patent-q-and-a"

# Quick calibration: 10 steps to measure starting reward
prime train run config.toml \
  --env-var OPENAI_API_KEY=sk-... \
  --override max_steps=10

prime train run config.toml --env-var OPENAI_API_KEY=sk-...