t5_translator/augment_data_abacus.py

import pandas as pd
import requests
import os
import time

# --- CONFIGURATION ---
# 1. Paste your NEW Abacus API Key here (Revoke the old one!)
ABACUS_API_KEY = "s2_0e301a7a1a524196a2cce70f72e620f0" 

# 2. Model Name (Double check this string in your Abacus UI if you get a 400 error)
MODEL_NAME = "claude-sonnet-4-20250514"

# 3. API Endpoint (Corrected for RouteLLM)
API_URL = "https://routellm.abacus.ai/v1/chat/completions"

# 4. File Paths
INPUT_CSV = "../../Data/The_SIGS_Lexicon_v1-STATIC_COPY-Master_Lexicon.csv"
OUTPUT_CSV = "../../Data/SIGS_v1_Expanded_Training_Data.csv"

def generate_sentences(token, definition, examples):
    """Asks AI to generate natural training sentences."""
    
    prompt = f"""
    I am building a dataset for an AI protocol called SIGS.
    
    TOKEN: "{token}"
    DEFINITION: "{definition}"
    EXISTING EXAMPLES: "{examples}"

    TASK:
    Generate 10 distinct, natural English sentences that a human user might say to an AI which strictly match this intent.
    
    RULES:
    1. Vary the tone (formal, casual, terse, polite, urgent).
    2. Do NOT use the token "{token}" in the output.
    3. Do NOT number the lines. Just provide the raw sentences separated by newlines.
    4. If the definition implies a specific action (like "Open file"), provide sentences that command that action.
    """

    payload = {
        "model": MODEL_NAME,
        "messages": [
            {"role": "system", "content": "You are a helpful data generation assistant."},
            {"role": "user", "content": prompt}
        ],
        "temperature": 0.8,
        "max_tokens": 500
    }
    
    headers = {
        "Authorization": f"Bearer {ABACUS_API_KEY}",
        "Content-Type": "application/json"
    }

    try:
        response = requests.post(API_URL, headers=headers, json=payload)
        
        # Error Handling
        if response.status_code != 200:
            print(f"⚠️ API Error {response.status_code} for {token}")
            print(f"Server Response: {response.text}") # Prints the actual error reason
            
            # Rate Limit Handling (429 = Too Many Requests)
            if response.status_code == 429:
                print("⏳ Rate limited. Sleeping for 10 seconds...")
                time.sleep(10)
            return []
            
        data = response.json()
        content = data['choices'][0]['message']['content']
        
        # Clean up output
        lines = [line.strip("- ").strip() for line in content.split('\n') if line.strip()]
        return lines

    except Exception as e:
        print(f"⚠️ Script Connection Error for {token}: {e}")
        return []

def main():
    if not os.path.exists(INPUT_CSV):
        print(f"❌ Input CSV not found at: {os.path.abspath(INPUT_CSV)}")
        return

    print(f"🚀 Connecting to Abacus AI ({API_URL})...")
    
    df = pd.read_csv(INPUT_CSV)
    total_rows = len(df)
    expanded_rows = []
    
    print(f"Processing {total_rows} tokens using {MODEL_NAME}...")

    for index, row in df.iterrows():
        token = row['wire_token']
        defn = row['definition']
        ex = str(row['examples'])
        
        print(f"[{index+1}/{total_rows}] Generating for {token}...")
        
        # 1. Keep original examples
        if ex and ex.lower() != 'nan':
            for e in ex.split('|'):
                expanded_rows.append({"input_text": e.strip(), "target_text": token})

        # 2. Generate SYNTHETIC examples
        new_sentences = generate_sentences(token, defn, ex)
        for s in new_sentences:
            expanded_rows.append({"input_text": s, "target_text": token})
            
        # Checkpoint every 20 rows
        if index % 20 == 0:
            pd.DataFrame(expanded_rows).to_csv(OUTPUT_CSV, index=False)
            print(f"   💾 Checkpoint saved ({len(expanded_rows)} pairs).")
            
        time.sleep(0.2) 

    # Final Save
    pd.DataFrame(expanded_rows).to_csv(OUTPUT_CSV, index=False)
    print(f"🎉 DONE! Saved to: {os.path.abspath(OUTPUT_CSV)}")

if __name__ == "__main__":
    main()