app.py

import gradio as gr
from setfit import SetFitModel
from sentence_transformers import SentenceTransformer, util
import re
import numpy as np
import traceback

# --- 1. SETUP & LOAD MODELS ---
print("Loading Models...")
try:
    # A. Load your Custom Trained Food Model (Bilingual)
    food_model = SetFitModel.from_pretrained("filter-search-model_v2")
    
    # B. Load Multilingual Vector Model (The "Bilingual Brain")
    # We use this instead of 'all-MiniLM-L6-v2' so it understands Vietnamese inputs
    vector_model = SentenceTransformer("JustscrAPIng/cultour-base-model-v2")
    
    print("Models Loaded Successfully!")
except Exception as e:
    print(f"Error loading models: {e}")

# --- 2. DEFINE DICTIONARIES ---

tag_definitions_en = {
    "religion": {
        'none': 'secular non-religious atheism no religion',
        'taoist': 'taoist taoism yin yang dao',
        'chinese_folk': 'chinese folk religion ancestor worship shenism',
        'buddhist': 'buddhist temple buddhism monk zen lotus',
        'christian': 'christian church jesus catholic protestant cross bible',
        'hinduism': 'hindu temple hinduism shiva vishnu',
        'vietnamese_folk': 'vietnamese folk religion mother goddess dao mau thanh mau',
        'khmer': 'khmer religion theravada buddhism cambodian style',
        'muslim': 'muslim islam mosque halal allah'
    },
    "arch_style": {
        'french': 'french colonial architecture indochina villa balcony yellow walls shutters',
        'modernism': 'modern architecture modernism brutalist concrete glass sleek minimal high-rise',
        'east asian': 'east asian oriental wooden structure curved roof',
        'contemporary': 'contemporary design new recent 21st century current',
        'baroque': 'baroque architecture ornate dramatic detailed grandiose european',
        'rococo': 'rococo style decorative pastel intricate playful',
        'vietnamese': 'traditional vietnamese architecture wooden red tile roof',
        'chinese': 'chinese architecture dragon courtyard feng shui',
        'romanesque': 'romanesque arches thick walls sturdy rounded',
        'gothic': 'gothic architecture pointed arches stained glass spires cathedral',
        'renaissance': 'renaissance symmetry domes columns classical proportion',
        'neo-romanesque': 'neo-romanesque revival 19th century style',
        'indochine': 'indochine style fusion french colonial and vietnamese tropical',
        'neo-baroque': 'neo-baroque revival grand opera house style',
        'hindu': 'hindu architecture cham style champa towers sandstone',
        'art deco': 'art deco retro geometric stylized 1920s 1930s',
        'khmer': 'khmer architecture angkor wat style stone temple prasat',
        'islamic': 'islamic architecture domes minarets geometric patterns',
        'flamboyant': 'flamboyant gothic flame-like intricate stone tracery',
        'beaux arts': 'beaux arts style monumental classical grand french school',
        'art nouveau': 'art nouveau organic lines flowery decorative curved'
    },
    "building_type": {
        'marketplace': 'marketplace market bazaar shopping buy food groceries ben thanh',
        'community_centre': 'community center public gathering cultural house hall',
        'courthouse': 'courthouse law legal judge justice government',
        'hospital': 'hospital clinic doctor medical emergency healthcare',
        'mortuary': 'mortuary funeral home dead burial services',
        'place_of_worship': 'place of worship temple church pagoda shrine mosque pray spiritual',
        'institution': 'institution school university college education academy library',
        'tomb': 'tomb grave mausoleum cemetery burial site',
        'heritage': 'heritage site historical building landmark monument preservation',
        'zoo': 'zoo animals botanical garden nature park',
        'museum': 'museum exhibition history art gallery display',
        'commercial': 'commercial building shop store mall office business trade',
        'government': 'government building administrative town hall committee',
        'residential': 'residential house home villa apartment living',
        'library': 'library books reading study archive',
        'fountain': 'fountain water feature square plaza',
        'theatre': 'theatre opera house performance stage cinema arts',
        'pagoda': 'pagoda tiered tower buddhist temple asian',
        'park': 'park green space garden nature relax trees',
        'reservoir': 'reservoir water lake dam utility',
        'restaurant': 'restaurant place to eat food dining hungry lunch dinner eatery',
        'cafe': 'cafe coffee shop drink sip tea chill work',
    },
    "food_type": {
         'vietnamese': '', 'international': '', 'asian': '', 'korean': '', 'grill': '', 
         'argentinian': '', 'japanese': '', 'lebanese': '', 'chinese': '', 'local food': '', 
         'pizza': '', 'mediterranean': '', 'thai': '', 'italian': '', 'hue style': '', 
         'indian': '', 'french': '', 'barbecue': '', 'banh cuon': '', 'nướng (grill)': '', 
         'sushi': '', 'beef bowl': '', 'noodles': '', 'curry': '', 'burger': '', 'seafood': '', 
         'american': '', 'pasta': '', 'thai hotpot': '', 'grilled fish': '', 'steamed fish': '', 
         'stir-fried vegetables': '', 'hotpot': '', 'sandwich': '', 'ice cream': '', 'crepe': '', 
         'coffee': '', 'fine dining': '', 'diner': '', 'steakhouse': '', 'mexican': '', 'wings': ''
    }
}

tag_definitions_vi = {
    "religion": {
        'none': 'không tôn giáo vô thần thế tục',
        'taoist': 'đạo giáo lão giáo âm dương',
        'chinese_folk': 'tín ngưỡng dân gian trung hoa thờ cúng tổ tiên',
        'buddhist': 'phật giáo chùa sư đạo phật nơi thờ tự người theo đạo phật hoa sen sư thầy thiền', 
        'christian': 'công giáo tin lành nhà thờ chúa giêsu thánh giá kinh thánh',
        'hinduism': 'ấn độ giáo đền thần shiva vishnu',
        'vietnamese_folk': 'tín ngưỡng dân gian việt nam đạo mẫu hầu đồng thánh mẫu',
        'khmer': 'phật giáo nam tông khmer chùa tháp',
        'muslim': 'hồi giáo thánh đường halal allah'
    },
    "arch_style": {
        'french': 'kiến trúc pháp thuộc địa đông dương biệt thự tường vàng cửa chớp',
        'modernism': 'kiến trúc hiện đại bê tông kính tối giản cao tầng tòa nhà kính',
        'east asian': 'kiến trúc đông á á đông gỗ mái cong',
        'contemporary': 'kiến trúc đương đại thiết kế mới thế kỷ 21',
        'baroque': 'kiến trúc baroque lộng lẫy chi tiết cầu kỳ châu âu',
        'rococo': 'phong cách rococo trang trí pastel tinh xảo',
        'vietnamese': 'kiến trúc truyền thống việt nam nhà gỗ mái ngói đỏ nhà rường',
        'chinese': 'kiến trúc trung hoa rồng sân trong phong thủy',
        'romanesque': 'kiến trúc romanesque vòm tròn tường dày',
        'gothic': 'kiến trúc gothic vòm nhọn kính màu tháp chuông nhà thờ',
        'renaissance': 'phục hưng mái vòm cột đối xứng cổ điển',
        'neo-romanesque': 'tân romanesque thế kỷ 19',
        'indochine': 'phong cách đông dương kết hợp pháp và việt nhiệt đới',
        'neo-baroque': 'tân baroque nhà hát lớn',
        'hindu': 'kiến trúc chăm pa tháp chăm gạch nung',
        'art deco': 'art deco hình học cách điệu thập niên 1920',
        'khmer': 'kiến trúc khmer angkor wat đền đá',
        'islamic': 'kiến trúc hồi giáo mái vòm tháp',
        'flamboyant': 'gothic rực cháy đá chạm khắc tinh xảo',
        'beaux arts': 'beaux arts hoành tráng cổ điển pháp',
        'art nouveau': 'tân nghệ thuật đường cong họa tiết hoa'
    },
    "building_type": {
        'marketplace': 'chợ bến thành mua sắm thực phẩm đồ khô',
        'community_centre': 'trung tâm cộng đồng nhà văn hóa hội trường',
        'courthouse': 'tòa án pháp luật xét xử',
        'hospital': 'bệnh viện phòng khám bác sĩ cấp cứu y tế',
        'mortuary': 'nhà tang lễ mai táng',
        'place_of_worship': 'nơi thờ tự chùa đền nhà thờ thánh đường miếu',
        'institution': 'trường học đại học thư viện giáo dục',
        'tomb': 'lăng mộ nghĩa trang',
        'heritage': 'di sản di tích lịch sử bảo tồn',
        'zoo': 'sở thú thảo cầm viên vườn bách thảo',
        'museum': 'bảo tàng triển lãm lịch sử nghệ thuật',
        'commercial': 'tòa nhà thương mại cửa hàng trung tâm mua sắm văn phòng',
        'government': 'cơ quan nhà nước ủy ban hành chính',
        'residential': 'nhà ở chung cư biệt thự căn hộ',
        'library': 'thư viện đọc sách tài liệu',
        'fountain': 'đài phun nước quảng trường',
        'theatre': 'nhà hát rạp chiếu phim sân khấu nghệ thuật',
        'pagoda': 'chùa tháp phật giáo',
        'park': 'công viên cây xanh vườn hoa',
        'reservoir': 'hồ chứa nước đập thủy điện',
        'restaurant': 'nhà hàng quán ăn tiệm cơm ăn uống',
        'cafe': 'quán cà phê trà sữa làm việc',
    }
}

# --- 3. AUTO-MERGE DICTIONARIES ---
print("Merging dictionaries...")
combined_definitions = {}
for category in ["religion", "arch_style", "building_type"]:
    combined_definitions[category] = {}
    for tag in tag_definitions_en[category]:
        en_desc = tag_definitions_en[category][tag]
        vi_desc = ""
        if category in tag_definitions_vi:
            vi_desc = tag_definitions_vi[category].get(tag, "")
        combined_desc = f"{en_desc} {vi_desc}".strip()
        combined_definitions[category][tag] = combined_desc

# --- 4. PRE-COMPUTE VECTORS ---
print("Pre-computing vectors...")
encoded_db = {}
for category in combined_definitions:
    tags_map = combined_definitions[category]
    encoded_db[category] = {
        'embeddings': vector_model.encode(list(tags_map.values()), convert_to_tensor=True),
        'names': list(tags_map.keys())
    }
print("Ready.")

# --- 5. EXTRACTION LOGIC (Strict & Patched) ---
def extract_tags(user_text):
    try:
        detected_data = {}
        text_lower = user_text.lower()
        
        # --- A. MANUAL OVERRIDES (Safety Net) ---
        # 1. Food
        if 'bún chả' in text_lower:
            detected_data['food_type'] = ['vietnamese', 'local food']
        elif 'bánh mì' in text_lower and 'pate' in text_lower:
            detected_data['food_type'] = ['sandwich', 'vietnamese']
        elif 'bánh mì' in text_lower:
            detected_data['food_type'] = ['vietnamese']
        
        # 2. Religion
        if 'đạo phật' in text_lower or 'chùa' in text_lower:
            detected_data['religion'] = ['buddhist']
        elif 'nhà thờ' in text_lower or 'công giáo' in text_lower:
            detected_data['religion'] = ['christian']
            
        # 3. Architecture
        if 'kính' in text_lower and 'cao tầng' in text_lower:
            detected_data['arch_style'] = ['modernism']
        elif 'gỗ' in text_lower and 'truyền thống' in text_lower:
            detected_data['arch_style'] = ['vietnamese']
        
        # --- B. AGE EXTRACTION ---
        match = re.search(r'\b(1[7-9]|20)\d{2}s?\b', user_text)
        detected_data['age'] = match.group(0).rstrip('s') if match else None

        # --- C. VECTOR SEARCH (Strict) ---
        user_embedding = vector_model.encode(user_text, convert_to_tensor=True)
        base_threshold = 0.29
        
        for category in ["religion", "arch_style", "building_type"]:
            if category in detected_data: continue # Skip overrides
            
            db_data = encoded_db[category]
            scores = util.cos_sim(user_embedding, db_data['embeddings'])[0]
            
            candidates = []
            for idx, score in enumerate(scores):
                if score > base_threshold:
                    candidates.append((db_data['names'][idx], score.item()))
            
            candidates.sort(key=lambda x: x[1], reverse=True)
            
            # STRICT LIMIT: Keep only top 2 to reduce noise
            candidates = candidates[:2]
            
            if candidates:
                best_score = candidates[0][1]
                relative_cutoff = best_score * 0.90 
                final_tags = [t[0] for t in candidates if t[1] >= relative_cutoff]
                detected_data[category] = final_tags
            else:
                detected_data[category] = None

        # --- D. SMART FOOD LOGIC (Your Old Strict Method) ---
        if 'food_type' not in detected_data:
            food_triggers = ['restaurant', 'cafe', 'street_food', 'marketplace']
            buildings = detected_data.get('building_type') or []
            is_food_venue = any(b in food_triggers for b in buildings)
            
            prediction_array = food_model.predict([user_text])
            prediction_clean = prediction_array.tolist()
            
            # 1. Validation List (From your old code)
            valid_food_keys = list(tag_definitions_en['food_type'].keys())
            
            # 2. Strict Filter
            filtered_food = [tag for tag in prediction_clean if tag in valid_food_keys]

            # 3. Assign
            if is_food_venue and filtered_food:
                detected_data['food_type'] = filtered_food
            else:
                detected_data['food_type'] = None

        return detected_data

    except Exception as e:
        return {"error": str(e), "traceback": traceback.format_exc()}

# --- 6. LAUNCH API ---
iface = gr.Interface(
    fn=extract_tags,
    inputs=gr.Textbox(label="User Search", placeholder="e.g. I want a french villa"),
    outputs=gr.JSON(label="Extracted Tags"),
    title="Cultour AI API (Bilingual)",
    description="Supports English and Vietnamese. Send queries to /api/predict"
)

iface.launch()