Revolutionizing Pharma AI with Afnio: How Auto-Optimized Agents Achieve a 45% Accuracy Boost in Minutes

$ pip install afnio

import os
import json
import re
from getpass import getpass
 
from bio_patent import BioPatent
from utils import print_html_table
 
import afnio
import afnio.tellurio as te
import afnio.cognitive as cog
import afnio.cognitive.functional as F
 
from afnio import set_backward_model_client, get_backward_model_client
from afnio.models.openai import AsyncOpenAI
from afnio.utils.data import DataLoader
from afnio.trainer import Trainer

if not (openai_api_key := os.getenv("OPENAI_API_KEY")):
    openai_api_key = getpass("🔑 Enter your OpenAI API key: ")
os.environ["OPENAI_API_KEY"] = openai_api_key
 
te.configure_logging("INFO")
te.login(api_key="TELLURIO_API_KEY")
run = te.init("dmpiergiacomo", "Therapeutic Extraction")

BATCH_SIZE = 25
 
training_data = BioPatent(split="train", root="data")
validation_data = BioPatent(split="val", root="data")
test_data = BioPatent(split="test", root="data")
 
train_dataloader = DataLoader(training_data, batch_size=BATCH_SIZE, shuffle=True)
val_dataloader = DataLoader(validation_data, batch_size=BATCH_SIZE, seed=42)
test_dataloader = DataLoader(test_data, batch_size=BATCH_SIZE, seed=42)

extractor_task = """Your task is to analyze the relevant paragraphs from a patent to determine what therapeutic target it addresses. First determine if it targets KRAS, and if so, how specifically. If it does not target KRAS, identify what other target it addresses.
 
INPUT:
You will receive:
- Relevant paragraphs from the patent
 
INSTRUCTIONS:
Based SOLELY on the provided paragraphs and context, classify this patent into EXACTLY ONE of these categories:
 
1. PAN_RAS_INHIBITOR: The patent describes compounds that inhibit multiple different RAS isoforms (KRAS, HRAS, NRAS) or multiple mutations across different RAS isoforms.
 
2. PAN_KRAS_INHIBITOR: The patent describes compounds that inhibit multiple different KRAS mutations (such as compounds effective against multiple mutations like G12C, G12D, G12V, etc.). The patent MUST explicitly state that it targets multiple KRAS mutations specifically, not just mention "RAS" or "K-RAS" in general.
 
3. MUTATION-SPECIFIC: The patent describes compounds targeting a specific KRAS mutation only (like only G12C or only G12D).
 
4. POTENTIAL_KRAS_INHIBITOR: The patent mentions RAS or KRAS in the title or abstract, but it's not 100% clear if it specifically targets KRAS. This includes patents that:
  - Mention RAS inhibitors without explicitly specifying KRAS
  - Describe a RAS inhibitor but don't clarify if it's a general RAS inhibitor or specifically targets KRAS
  - Where the targeting of KRAS is ambiguous
  - Where the specificity between general RAS and KRAS is not clearly stated
  - Use ambiguous phrases like "Ras protein-related" or "K-Ras protein-related" without clear specificity
  - Mention "RAS" or "K-RAS" without explicitly stating which mutations or isoforms are targeted
 
5. OTHER_TARGET: The patent does NOT target KRAS but instead targets another therapeutic target. This could be any protein, enzyme, receptor, pathway, or biological process (not limited to EGFR, BRAF, MEK, etc.). Be specific about what the main therapeutic target is.
 
IMPORTANT GUIDELINES:
- Focus on paragraphs that show the actual invention, not background information
- Background information mentioning KRAS is not sufficient for KRAS classification; the compound must actively target KRAS
- Pay special attention to "method of treatment" claims or paragraphs that specifically mention target proteins or mechanisms
- Any paragraph showing the invention or compounds can treat multiple different RAS isoforms indicates PAN_RAS_INHIBITOR
- For PAN_KRAS_INHIBITOR, the patent MUST explicitly state that it targets multiple KRAS mutations specifically
- Mere mentions of "RAS" or "K-RAS" without specific mutation information should be classified as POTENTIAL_KRAS_INHIBITOR
- Phrases like "Ras protein-related" or "K-Ras protein-related" without clear specificity should be classified as POTENTIAL_KRAS_INHIBITOR
- Look for method claims, examples, or embodiments that specifically mention targets
- If no KRAS targeting is mentioned but other targets are discussed as the primary target, classify as OTHER_TARGET
- If the patent targets a general mechanism or pathway rather than a specific protein, name that mechanism/pathway in the target field
- The presence of multiple KRAS mutations in the text does NOT automatically mean it's a pan-KRAS inhibitor. Look for explicit statements about treating multiple mutations
- If the patent mentions RAS inhibitors but doesn't explicitly specify KRAS, classify as POTENTIAL_KRAS_INHIBITOR
"""
 
extractor_out_format = """
Return your answer in this exact JSON format:
{{
"classification": "PAN_RAS_INHIBITOR" or "PAN_KRAS_INHIBITOR" or "<MUTATION>_KRAS_INHIBITOR" or "POTENTIAL_KRAS_INHIBITOR" or "OTHER_TARGET",
"target": "RAS" or "KRAS" or the specific other target (e.g., "EGFR", "BRAF", "Tyrosine Kinase", "PI3K pathway", etc.),
"explanation": "Your detailed reasoning based on the provided paragraphs"
}}
 
The output must be valid JSON that can be parsed by json.loads().
Include specific evidence from the paragraphs.
"""
 
extractor_template = """
Relevant paragraphs from the patent: {relevant_content}
"""
 
EXTRACTOR_RESPONSE_FORMAT = {
    "type": "json_schema",
    "json_schema": {
        "name": "classification_response_schema",
        "schema": {
            "type": "object",
            "properties": {
                "classification": {"type": "string"},
                "target": {"type": "boolean"},
                "explanation": {"type": "string"},
            },
            "required": ["classification", "target", "explanation"],
        },
    },
}

set_backward_model_client("openai/gpt-4o")
fw_model_client = AsyncOpenAI()
optim_model_client = AsyncOpenAI()

class TherapeuticExtractor(cog.Module):
 
  def __init__(self):
    super().__init__()
    self.extractor_task = cog.Parameter(data=extractor_task, role="system prompt of classification task", requires_grad=True)
    self.extractor_format = cog.Parameter(data=extractor_out_format, role="classification output format", requires_grad=False)
    self.extractor_template = afnio.Variable(data=extractor_template, role="input template to classifier")
    self.extractor_chat = cog.ChatCompletion()
 
  def forward(self, fwd_model, inputs, **completion_args):
    classification_messages = [
      {"role": "system", "content": [self.extractor_task, self.extractor_format]},
      {"role": "user", "content": [self.extractor_template]},
    ]
    return self.extractor_chat(fwd_model, classification_messages, inputs=inputs, response_format=EXTRACTOR_RESPONSE_FORMAT, **completion_args)
  
  def training_step(self, batch, batch_idx):
    X, y = batch
    inputs = {"relevant_content": X}
    pred = self(fw_model_client, inputs=inputs, model="gpt-4o-mini", temperature=0.0)
    pred.data = [json.loads(re.sub(r"^\`\`\`json\\n|\\n\`\`\`$", "", item))["classification"] for item in pred.data]
    loss = F.exact_match_evaluator(pred, y)
    return {"loss": loss, "accuracy": loss[0].data / len(y.data)}
 
  def validation_step(self, batch, batch_idx):
    return self.training_step(batch, batch_idx)
 
  def test_step(self, batch, batch_idx):
    return self.validation_step(batch, batch_idx)
 
  def configure_optimizers(self):
    optimizer = afnio.optim.TGD(self.parameters(), model_client=optim_model_client, momentum=3, model="gpt-4o")
    return optimizer

agent = TherapeuticExtractor()
trainer = Trainer(max_epochs=6, enable_agent_summary=False)

llm_clients=[fw_model_client, get_backward_model_client(), optim_model_client]
 
trainer.test(agent=agent, test_dataloader=test_dataloader, llm_clients=llm_clients)
trainer.fit(agent=agent, train_dataloader=train_dataloader, val_dataloader=val_dataloader, llm_clients=llm_clients)

checkpoint_name = "checkpoints/checkpoint_epoch3_20250803-181033.hf"
checkpoint = afnio.load(checkpoint_name)
best_agent = TherapeuticExtractor()
best_agent.load_state_dict(checkpoint['agent_state_dict'], model_clients={"extractor_chat.forward_model_client": fw_model_client})
 
trainer.test(agent=best_agent, test_dataloader=test_dataloader, llm_clients=llm_clients, step=6)
 
run.finish()