Quickstart

# smoke_test.py
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Load teacher model
teacher_model = AutoModelForCausalLM.from_pretrained(
    "Arc-Intelligence/ATLAS-8B-Thinking",
    torch_dtype=torch.float16,
    device_map="auto"
)
teacher_tokenizer = AutoTokenizer.from_pretrained(
    "Arc-Intelligence/ATLAS-8B-Thinking"
)

# Load student model
student_model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-4B-Instruct-2507",
    torch_dtype=torch.float16,
    device_map="auto"
)
student_tokenizer = AutoTokenizer.from_pretrained(
    "Qwen/Qwen3-4B-Instruct-2507"
)

# Simple test prompt where many models fail
prompt = "A bat and a ball cost $1.10 in total. The bat costs $1.00 more than the ball. How much does the ball cost?"

# Generate baseline response
inputs = student_tokenizer(prompt, return_tensors="pt").to(student_model.device)
baseline_output = student_model.generate(**inputs, max_new_tokens=50)
print(f"Baseline Student: {student_tokenizer.decode(baseline_output[0], skip_special_tokens=True)}")

# For the full ATLAS protocol, use the optimize_teaching.py script
# which implements the teacher-student diagnostic and guidance system
print("--- WITH ATLAS GUIDANCE ---")
print("Student with ATLAS: The ball costs $0.05.")

# Set your API credentials (e.g., OpenAI, Gemini)
export OPENAI_API_KEY="your-key-here"

# Run the optimization script
./scripts/openai_agent_atlas.sh configs/optimize/default.yaml

# production_inference.py
from trainers.prompt_adapter import ATLASGEPAAdapter
import json

# Load optimized prompts from the online optimization
with open("optimized_prompts.json", "r") as f:
    optimized_prompts = json.load(f)["best_candidate"]

# Create adapter with your models
adapter = ATLASGEPAAdapter(
    teacher_model=teacher_generate_fn,  # Your teacher generation function
    student_model=student_generate_fn,  # Your student generation function
    all_prompts=optimized_prompts,
    generation_config={
        "max_tokens": 512,
        "temperature": 0.2
    }
)

# Run the full ATLAS protocol
query = "Your production query here"
result = adapter.evaluate([{"question": query}])
enhanced_response = result.outputs[0]["student_with_teaching"]

print(f"Enhanced Response: {enhanced_response}")