scripts/fix_ner_samples.py

"""Fix the 17 new NER fixture samples added in Wave 3:
- Add proper Vietnamese diacritics to time_preference samples
- Fix gold annotations to be consistent with model keyword dictionaries
- Replace Bali Mount Batur with simpler text (avoid 'batu' FP)
- Fix 'Family of 4' sample to not conflict with travel_style suppression logic
"""
import json
import os

DATA_PATH = os.path.join(os.path.dirname(__file__), '..', 'app', 'data', 'datasets', 'ner_eval.json')

CORRECTED = [
    # --- time_preference samples (12) with proper diacritics ---
    # NOTE: each destination city annotated as BOTH 'location' AND 'destination'
    # (matching convention in original 85 samples; model returns both)
    {
        'text': 'Muốn đi tham quan Hạ Long buổi sáng sớm trước khi tàu đông khách',
        'language': 'vi',
        'entities': [
            {'type': 'location', 'value': 'ha-long'},
            {'type': 'destination', 'value': 'ha-long'},
            {'type': 'time_preference', 'value': 'early_morning'},
        ]
    },
    {
        'text': 'Hội An 3 ngày, đi phố cổ buổi chiều và xem đèn lồng buổi tối',
        'language': 'vi',
        'entities': [
            {'type': 'location', 'value': 'hoi-an'},
            {'type': 'destination', 'value': 'hoi-an'},
            {'type': 'duration', 'value': {'days': 3, 'nights': 2}},
            {'type': 'time_preference', 'value': 'afternoon'},
            {'type': 'time_preference', 'value': 'evening'},
        ]
    },
    {
        'text': 'Chụp ảnh hoàng hôn ở Mũi Né, thích đi vào buổi chiều',
        'language': 'vi',
        'entities': [
            {'type': 'location', 'value': 'mui-ne'},
            {'type': 'destination', 'value': 'mui-ne'},
            {'type': 'activity_preference', 'value': 'photography'},
            {'type': 'time_preference', 'value': 'evening'},
        ]
    },
    {
        'text': 'Sapa trekking buổi sáng, nghỉ trưa rồi tiếp tục leo núi',
        'language': 'vi',
        'entities': [
            {'type': 'location', 'value': 'sapa'},
            {'type': 'destination', 'value': 'sapa'},
            {'type': 'activity_preference', 'value': 'trekking'},
            {'type': 'time_preference', 'value': 'morning'},
            {'type': 'time_preference', 'value': 'noon'},
        ]
    },
    {
        'text': 'Hà Nội 2 ngày 1 đêm, tham quan buổi sáng, mua sắm buổi chiều',
        'language': 'vi',
        'entities': [
            {'type': 'location', 'value': 'hanoi'},
            {'type': 'destination', 'value': 'hanoi'},
            {'type': 'duration', 'value': {'days': 2, 'nights': 1}},
            {'type': 'activity_preference', 'value': 'sightseeing'},
            {'type': 'activity_preference', 'value': 'shopping'},
            {'type': 'time_preference', 'value': 'morning'},
            {'type': 'time_preference', 'value': 'afternoon'},
        ]
    },
    {
        'text': 'Nha Trang 4 ngày, lặn biển buổi sáng sớm và đi chợ đêm buổi tối',
        'language': 'vi',
        'entities': [
            {'type': 'location', 'value': 'nha-trang'},
            {'type': 'destination', 'value': 'nha-trang'},
            {'type': 'duration', 'value': {'days': 4, 'nights': 3}},
            {'type': 'activity_preference', 'value': 'diving'},
            {'type': 'time_preference', 'value': 'early_morning'},
            {'type': 'time_preference', 'value': 'evening'},
        ]
    },
    {
        'text': 'Đi Phú Quốc ăn hải sản buổi trưa, ngắm hoàng hôn buổi chiều',
        'language': 'vi',
        'entities': [
            {'type': 'location', 'value': 'phu-quoc'},
            {'type': 'destination', 'value': 'phu-quoc'},
            {'type': 'cuisine', 'value': 'seafood'},
            {'type': 'time_preference', 'value': 'noon'},
            {'type': 'time_preference', 'value': 'evening'},
        ]
    },
    {
        # Bangkok: 'night markets in the evening' -> evening only (night compound suppressed)
        'text': 'I want to explore Bangkok temples in the morning and night markets in the evening',
        'language': 'en',
        'entities': [
            {'type': 'location', 'value': 'bangkok'},
            {'type': 'destination', 'value': 'bangkok'},
            {'type': 'activity_preference', 'value': 'sightseeing'},
            {'type': 'time_preference', 'value': 'morning'},
            {'type': 'time_preference', 'value': 'evening'},
        ]
    },
    {
        'text': 'Kyoto 5 days, visit shrines early morning before crowds arrive',
        'language': 'en',
        'entities': [
            {'type': 'location', 'value': 'kyoto'},
            {'type': 'destination', 'value': 'kyoto'},
            {'type': 'duration', 'value': {'days': 5, 'nights': 4}},
            {'type': 'activity_preference', 'value': 'sightseeing'},
            {'type': 'time_preference', 'value': 'early_morning'},
        ]
    },
    {
        'text': 'Singapore 3 nights, food tour at noon and rooftop bar in the evening',
        'language': 'en',
        'entities': [
            {'type': 'location', 'value': 'singapore'},
            {'type': 'destination', 'value': 'singapore'},
            {'type': 'duration', 'value': {'days': 3, 'nights': 3}},
            {'type': 'cuisine', 'value': 'local_specialty'},
            {'type': 'time_preference', 'value': 'noon'},
            {'type': 'time_preference', 'value': 'evening'},
        ]
    },
    {
        # Removed 'Mount Batur' to avoid 'batu' location FP; simplified to Bali morning hike
        'text': 'Bali morning hike, afternoon spa and massage session',
        'language': 'en',
        'entities': [
            {'type': 'location', 'value': 'bali'},
            {'type': 'destination', 'value': 'bali'},
            {'type': 'activity_preference', 'value': 'trekking'},
            {'type': 'activity_preference', 'value': 'spa'},
            {'type': 'time_preference', 'value': 'morning'},
            {'type': 'time_preference', 'value': 'afternoon'},
        ]
    },
    {
        'text': 'Paris night tour: Eiffel Tower at night, brunch in the morning',
        'language': 'en',
        'entities': [
            {'type': 'location', 'value': 'paris'},
            {'type': 'destination', 'value': 'paris'},
            {'type': 'activity_preference', 'value': 'sightseeing'},
            {'type': 'time_preference', 'value': 'night'},
            {'type': 'time_preference', 'value': 'morning'},
        ]
    },
    # --- people edge cases (5 samples) ---
    {
        # 'Solo backpacker 7 days' - explicit '7 days' duration
        # gold: backpacker+adventure (both fired by model keywords)
        'text': 'Solo backpacker trip to Sapa 7 days, hiking and photography',
        'language': 'en',
        'entities': [
            {'type': 'location', 'value': 'sapa'},
            {'type': 'destination', 'value': 'sapa'},
            {'type': 'duration', 'value': {'days': 7, 'nights': 6}},
            {'type': 'people', 'value': {'adults': 1, 'children': 0, 'infants': 0}},
            {'type': 'travel_style', 'value': 'backpacker'},
            {'type': 'travel_style', 'value': 'adventure'},
            {'type': 'activity_preference', 'value': 'trekking'},
            {'type': 'activity_preference', 'value': 'photography'},
        ]
    },
    {
        'text': 'Couple trip to Bali 7 days, honeymoon package',
        'language': 'en',
        'entities': [
            {'type': 'location', 'value': 'bali'},
            {'type': 'destination', 'value': 'bali'},
            {'type': 'duration', 'value': {'days': 7, 'nights': 6}},
            {'type': 'people', 'value': {'adults': 2, 'children': 0, 'infants': 0}},
            {'type': 'travel_style', 'value': 'honeymoon'},
        ]
    },
    {
        'text': 'Group of 8 friends to Da Nang beach, budget hostel',
        'language': 'en',
        'entities': [
            {'type': 'location', 'value': 'da-nang'},
            {'type': 'destination', 'value': 'da-nang'},
            {'type': 'people', 'value': {'adults': 8, 'children': 0, 'infants': 0}},
            {'type': 'hotel_type', 'value': 'hostel'},
        ]
    },
    {
        'text': 'Đi Phú Quốc 2 vợ chồng 5 ngày nghỉ dưỡng',
        'language': 'vi',
        'entities': [
            {'type': 'location', 'value': 'phu-quoc'},
            {'type': 'destination', 'value': 'phu-quoc'},
            {'type': 'duration', 'value': {'days': 5, 'nights': 4}},
            {'type': 'people', 'value': {'adults': 2, 'children': 0, 'infants': 0}},
            {'type': 'travel_style', 'value': 'relaxation'},
        ]
    },
    {
        'text': 'Gia đình 5 người đi Đà Lạt 4 ngày, có 2 trẻ em',
        'language': 'vi',
        'entities': [
            {'type': 'location', 'value': 'da-lat'},
            {'type': 'destination', 'value': 'da-lat'},
            {'type': 'duration', 'value': {'days': 4, 'nights': 3}},
            {'type': 'people', 'value': {'adults': 3, 'children': 2, 'infants': 0}},
            {'type': 'travel_style', 'value': 'family'},
        ]
    },
]

assert len(CORRECTED) == 17, f'Expected 17, got {len(CORRECTED)}'

data = json.load(open(DATA_PATH, encoding='utf-8'))
print(f'Before: {len(data)} samples')
data[85:] = CORRECTED
print(f'After: {len(data)} samples')

tp = sum(1 for s in data if any(e['type'] == 'time_preference' for e in s.get('entities', [])))
print(f'time_preference samples: {tp}')

json.dump(data, open(DATA_PATH, 'w', encoding='utf-8'), ensure_ascii=False, indent=2)
print('Saved.')