generate_training_data.py

"""
Generate training data from sample_data.json
Procurement Assistant - product recommendations, pricing, supplier info
"""

import json
import random
from collections import defaultdict

def load_data(filename="sample_data.json"):
    """Load procurement data."""
    with open(filename, 'r', encoding='utf-8') as f:
        return json.load(f)

# Question templates - English
BUY_PRODUCT_QUESTIONS = [
    "I want to buy {product}",
    "I need to purchase {product}",
    "Looking for {product}",
    "Where can I get {product}?",
    "Need {product}, any recommendations?",
    "Suggest supplier for {product}",
    "Who sells {product}?",
    "Find me {product}",
]

PRICE_QUESTIONS = [
    "What is the price for {product}?",
    "How much does {product} cost?",
    "Price of {product}?",
    "Cost for {product}?",
    "{product} price?",
]

SUPPLIER_QUESTIONS = [
    "What does {supplier} sell?",
    "Products from {supplier}?",
    "Tell me about {supplier}",
    "What can I buy from {supplier}?",
    "{supplier} product list",
]

COMPARE_QUESTIONS = [
    "Compare suppliers for {product}",
    "Who else sells {product}?",
    "Other suppliers for {product}?",
    "Alternatives for {product}?",
    "Different vendors for {product}?",
]

QUANTITY_QUESTIONS = [
    "I need {qty} {unit} of {product}",
    "Order {qty} {unit} of {product}",
    "Buy {qty} {product}",
    "Quote for {qty} {unit} {product}",
]

def generate_samples(data):
    """Generate training samples."""
    samples = []

    # Build indexes
    supplier_products = defaultdict(list)  # supplier -> [products]
    product_suppliers = defaultdict(list)  # product_name -> [suppliers with prices]
    all_products = []

    for record in data:
        contractor = record.get('contractor')
        if not contractor or not contractor.strip():
            continue
        contractor = contractor.strip()

        products = record.get('products', [])
        record_usd = record.get('usd_amount')
        total_amount = record.get('total_amount')
        notes = record.get('notes', '')

        for prod in products:
            name = prod.get('name', '').strip()
            if not name:
                continue

            param = str(prod.get('parameter', '') or '').strip()
            unit = prod.get('unit', 'шт')
            qty = prod.get('quantity', 0)
            unit_price = prod.get('unit_price')
            amount = prod.get('amount')
            prod_usd = prod.get('usd_amount') or record_usd

            product_info = {
                'name': name,
                'param': param,
                'full_name': f"{name} ({param})" if param else name,
                'unit': unit,
                'qty': qty,
                'unit_price': unit_price,
                'amount': amount,
                'usd': prod_usd,
                'supplier': contractor,
                'notes': notes,
            }

            supplier_products[contractor].append(product_info)
            product_suppliers[name.lower()].append(product_info)
            all_products.append(product_info)

    print(f"Indexed: {len(supplier_products)} suppliers, {len(all_products)} products")

    # 1. "I want to buy X" → Recommendation
    for p in all_products:
        question = random.choice(BUY_PRODUCT_QUESTIONS).format(product=p['full_name'])

        answer = f"I recommend purchasing **{p['full_name']}** from:\n\n"
        answer += f"**Supplier:** {p['supplier']}\n"

        if p['unit_price'] and p['unit_price'] > 0:
            answer += f"**Unit Price:** {p['unit_price']:,.0f} UZS per {p['unit']}\n"
            if p['amount'] and p['amount'] > 0:
                answer += f"**Total ({p['qty']} {p['unit']}):** {p['amount']:,.0f} UZS\n"
        elif p['usd']:
            answer += f"**Price:** ${p['usd']:,.2f} USD\n"

        # Other products from same supplier
        other_prods = [x for x in supplier_products[p['supplier']] if x['name'] != p['name']][:3]
        if other_prods:
            answer += f"\n**Other products from {p['supplier']}:**\n"
            for op in other_prods:
                if op['unit_price'] and op['unit_price'] > 0:
                    answer += f"- {op['full_name']}: {op['unit_price']:,.0f} UZS/{op['unit']}\n"
                elif op['usd']:
                    answer += f"- {op['full_name']}: ${op['usd']:,.2f} USD\n"
                else:
                    answer += f"- {op['full_name']}\n"

        # Alternative suppliers
        alt_suppliers = [x for x in product_suppliers[p['name'].lower()] if x['supplier'] != p['supplier']][:2]
        if alt_suppliers:
            answer += f"\n**Alternative suppliers:**\n"
            for alt in alt_suppliers:
                if alt['unit_price'] and alt['unit_price'] > 0:
                    answer += f"- {alt['supplier']}: {alt['unit_price']:,.0f} UZS/{alt['unit']}\n"
                elif alt['usd']:
                    answer += f"- {alt['supplier']}: ${alt['usd']:,.2f} USD\n"
                else:
                    answer += f"- {alt['supplier']}\n"

        samples.append({"instruction": question, "input": "", "output": answer.strip()})

    # 2. Price questions
    for p in all_products:
        if not p['unit_price'] and not p['usd']:
            continue

        question = random.choice(PRICE_QUESTIONS).format(product=p['full_name'])

        answer = f"**{p['full_name']}** pricing:\n\n"
        answer += f"**Supplier:** {p['supplier']}\n"

        if p['unit_price'] and p['unit_price'] > 0:
            answer += f"**Unit Price:** {p['unit_price']:,.0f} UZS per {p['unit']}\n"
            if p['qty'] and p['amount']:
                answer += f"**Quantity:** {p['qty']} {p['unit']}\n"
                answer += f"**Total Amount:** {p['amount']:,.0f} UZS\n"

        if p['usd']:
            answer += f"**USD Price:** ${p['usd']:,.2f}\n"

        samples.append({"instruction": question, "input": "", "output": answer.strip()})

    # 3. Supplier info questions
    for supplier, prods in supplier_products.items():
        question = random.choice(SUPPLIER_QUESTIONS).format(supplier=supplier)

        answer = f"**{supplier}** offers:\n\n"

        for i, p in enumerate(prods[:8], 1):
            answer += f"{i}. {p['full_name']}\n"
            if p['unit_price'] and p['unit_price'] > 0:
                answer += f"   Price: {p['unit_price']:,.0f} UZS/{p['unit']}\n"
            elif p['usd']:
                answer += f"   Price: ${p['usd']:,.2f} USD\n"

        if len(prods) > 8:
            answer += f"\n...and {len(prods) - 8} more products."

        samples.append({"instruction": question, "input": "", "output": answer.strip()})

    # 4. Compare/alternatives questions
    for product_name, suppliers in product_suppliers.items():
        if len(suppliers) < 2:
            continue

        first_prod = suppliers[0]
        question = random.choice(COMPARE_QUESTIONS).format(product=first_prod['name'])

        answer = f"Suppliers for **{first_prod['name']}**:\n\n"

        seen = set()
        count = 0
        for s in suppliers:
            if s['supplier'] in seen:
                continue
            seen.add(s['supplier'])
            count += 1

            answer += f"{count}. **{s['supplier']}**\n"
            answer += f"   - {s['full_name']}\n"
            if s['unit_price'] and s['unit_price'] > 0:
                answer += f"   - Price: {s['unit_price']:,.0f} UZS/{s['unit']}\n"
            elif s['usd']:
                answer += f"   - Price: ${s['usd']:,.2f} USD\n"

            if count >= 5:
                break

        samples.append({"instruction": question, "input": "", "output": answer.strip()})

    # 5. Quantity questions
    for p in all_products:
        if not p['qty'] or p['qty'] <= 0:
            continue
        if not p['unit_price'] and not p['usd']:
            continue

        question = random.choice(QUANTITY_QUESTIONS).format(
            qty=int(p['qty']), unit=p['unit'], product=p['name']
        )

        answer = f"For **{int(p['qty'])} {p['unit']}** of **{p['full_name']}**:\n\n"
        answer += f"**Recommended Supplier:** {p['supplier']}\n"

        if p['unit_price'] and p['unit_price'] > 0:
            answer += f"**Unit Price:** {p['unit_price']:,.0f} UZS\n"
            if p['amount']:
                answer += f"**Total Cost:** {p['amount']:,.0f} UZS\n"
        elif p['usd']:
            answer += f"**Price:** ${p['usd']:,.2f} USD\n"

        # Alternatives
        alts = [x for x in product_suppliers[p['name'].lower()] if x['supplier'] != p['supplier']][:2]
        if alts:
            answer += "\n**Alternatives:**\n"
            for alt in alts:
                if alt['unit_price']:
                    answer += f"- {alt['supplier']}: {alt['unit_price']:,.0f} UZS/{alt['unit']}\n"
                elif alt['usd']:
                    answer += f"- {alt['supplier']}: ${alt['usd']:,.2f} USD\n"

        samples.append({"instruction": question, "input": "", "output": answer.strip()})

    random.shuffle(samples)
    return samples

def save_jsonl(samples, filename="train_data.jsonl"):
    """Save to JSONL."""
    with open(filename, 'w', encoding='utf-8') as f:
        for s in samples:
            json.dump(s, f, ensure_ascii=False)
            f.write('\n')
    print(f"Saved {len(samples)} samples to {filename}")

def main():
    print("Loading sample_data.json...")
    data = load_data("sample_data.json")
    print(f"Loaded {len(data)} records")

    print("\nGenerating training samples...")
    samples = generate_samples(data)

    print(f"\nTotal samples: {len(samples)}")

    # Preview
    print("\n" + "="*50)
    print("SAMPLE PREVIEW")
    print("="*50)
    for i, s in enumerate(random.sample(samples, min(3, len(samples)))):
        print(f"\n[Sample {i+1}]")
        print(f"Q: {s['instruction']}")
        print(f"A: {s['output']}")
        print("-"*40)

    save_jsonl(samples)
    print("\nDone!")

if __name__ == "__main__":
    main()