Running Benchmarks

from chemlactica.mol_opt.metrics import top_auc

# Calculate Top-100 AUC
auc_score = top_auc(
    buffer=oracle.mol_buffer,
    top_n=100,
    finish=oracle.finish,
    freq_log=100,
    max_oracle_calls=1000
)

def log_intermediate(self):
    scores = [v[0] for v in self.mol_buffer.values()][-self.max_oracle_calls:]
    scores_sorted = sorted(scores, reverse=True)[:100]
    n_calls = len(self.mol_buffer)

    score_avg_top1 = np.max(scores_sorted)
    score_avg_top10 = np.mean(scores_sorted[:10])
    score_avg_top100 = np.mean(scores_sorted)

    print(f"{n_calls}/{self.max_oracle_calls} | ",
          f'avg_top1: {score_avg_top1:.3f} | '
          f'avg_top10: {score_avg_top10:.3f} | '
          f'avg_top100: {score_avg_top100:.3f}')

100/1000 | avg_top1: 156.432 | avg_top10: 142.156 | avg_top100: 128.743
200/1000 | avg_top1: 178.921 | avg_top10: 165.234 | avg_top100: 151.892

from chemlactica.mol_opt.metrics import internal_diversity
from rdkit.Chem import AllChem
import numpy as np

# Generate fingerprints
fingerprints = np.array([
    AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits=2048)
    for mol in molecules
])

# Calculate diversity
diversity = internal_diversity(
    molecule_fingerprints=fingerprints,
    device='cpu',
    fp_type='morgan',
    p=1
)

print(f"Internal diversity: {diversity:.3f}")

from chemlactica.mol_opt.metrics import average_agg_tanimoto

# Compare generated molecules to training set
similarity = average_agg_tanimoto(
    stock_vecs=reference_fingerprints,
    gen_vecs=generated_fingerprints,
    batch_size=5000,
    agg='max',  # or 'mean'
    device='cuda',
    p=1
)

print(f"Average max Tanimoto similarity: {similarity:.3f}")

import yaml
import os
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from chemlactica.mol_opt.optimization import optimize
from chemlactica.mol_opt.utils import set_seed

# Load config
config = yaml.safe_load(open("config.yaml"))

# Load model once
model = AutoModelForCausalLM.from_pretrained(
    config["checkpoint_path"],
    torch_dtype=torch.bfloat16
).to(config["device"])

tokenizer = AutoTokenizer.from_pretrained(
    config["tokenizer_path"],
    padding_side="left"
)

# Run multiple trials
seeds = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29]
results = []

for i, seed in enumerate(seeds):
    print(f"\n=== Trial {i+1}/{len(seeds)} (seed={seed}) ===")
    
    set_seed(seed)
    oracle = YourOracle(max_oracle_calls=1000)
    
    config["log_dir"] = os.path.join(
        output_dir,
        f"results_trial_{i}_seed_{seed}.log"
    )
    config["max_possible_oracle_score"] = oracle.max_possible_oracle_score
    
    optimize(model, tokenizer, oracle, config)
    
    # Collect results
    results.append({
        'seed': seed,
        'top1': np.max([v[0] for v in oracle.mol_buffer.values()]),
        'top10': np.mean(sorted([v[0] for v in oracle.mol_buffer.values()], reverse=True)[:10]),
        'top100': np.mean(sorted([v[0] for v in oracle.mol_buffer.values()], reverse=True)[:100]),
        'num_molecules': len(oracle.mol_buffer)
    })

# Aggregate results
import pandas as pd

df = pd.DataFrame(results)
print("\n=== Aggregate Results ===")
print(df.describe())
df.to_csv(os.path.join(output_dir, "benchmark_summary.csv"), index=False)

import re
import pandas as pd

def parse_log_file(log_path):
    molecules = []
    
    with open(log_path, 'r') as f:
        for line in f:
            if line.startswith('generated smiles:'):
                match = re.search(r'generated smiles: (.+?), score: ([\d.]+)', line)
                if match:
                    smiles, score = match.groups()
                    molecules.append({
                        'smiles': smiles,
                        'score': float(score)
                    })
    
    return pd.DataFrame(molecules)

# Parse and analyze
df = parse_log_file('results.log')
print(f"Total molecules: {len(df)}")
print(f"Unique molecules: {df['smiles'].nunique()}")
print(f"Top-10 average: {df.nlargest(10, 'score')['score'].mean():.3f}")

import matplotlib.pyplot as plt
import seaborn as sns

# Score distribution
plt.figure(figsize=(10, 6))
sns.histplot(data=df, x='score', bins=50)
plt.title('Distribution of Oracle Scores')
plt.xlabel('Score')
plt.ylabel('Count')
plt.savefig('score_distribution.png')

# Top scores over time
df_sorted = df.sort_values('score', ascending=False).reset_index(drop=True)
top_scores = df_sorted.head(100)['score'].values

plt.figure(figsize=(10, 6))
plt.plot(range(1, 101), top_scores, marker='o')
plt.title('Top 100 Molecules by Score')
plt.xlabel('Rank')
plt.ylabel('Score')
plt.grid(True, alpha=0.3)
plt.savefig('top_scores.png')

import pandas as pd
import matplotlib.pyplot as plt

# Load results from multiple models
model_results = {
    'ChemLactica-125M': pd.read_csv('results/125m_summary.csv'),
    'ChemLactica-1B': pd.read_csv('results/1b_summary.csv'),
    'Baseline': pd.read_csv('results/baseline_summary.csv')
}

# Compare top-10 scores
comparison = pd.DataFrame({
    model: data['top10'].describe()[['mean', 'std', 'min', 'max']]
    for model, data in model_results.items()
}).T

print("\n=== Top-10 Score Comparison ===")
print(comparison)

# Box plot comparison
fig, ax = plt.subplots(figsize=(10, 6))
data_to_plot = [data['top10'].values for data in model_results.values()]
ax.boxplot(data_to_plot, labels=model_results.keys())
ax.set_ylabel('Top-10 Average Score')
ax.set_title('Model Performance Comparison')
plt.xticks(rotation=45)
plt.tight_layout()
plt.savefig('model_comparison.png')