dataset-web.py

import os
import json
import random
import gradio as gr

# Paths
JSON_DIR = "/research/iprobe-shuklan3/face-DPO/conversations"
BASE_IMAGE_PATH = "/research/iprobe-shuklan3/face-DPO/Dataset/Train"

# Load all JSON files
def list_json_files():
    return [f for f in os.listdir(JSON_DIR) if f.endswith(".json")]

# Load JSON content
def load_json(file_name):
    with open(os.path.join(JSON_DIR, file_name), "r") as f:
        return json.load(f)

# Select a random entry from JSON
def random_entry(json_data):
    return random.choice(json_data["conversations"]) if "conversations" in json_data else None

# Gradio interface function
def show_random_entry(selected_json):
    if not selected_json:
        return "Select a JSON file first", [], "", ""
    
    json_data = load_json(selected_json)
    entry = random_entry(json_data)
    if not entry:
        return "No entries found in JSON.", [], "", ""
    
    # Extract info
    human_prompt = entry.get("conversations", [{}])[0].get("value", "")
    chosen_text = entry.get("chosen", {}).get("value", "")
    rejected_text = entry.get("rejected", {}).get("value", "")
    
    # Build image paths
    images = entry.get("images", [])
    img_paths = []
    for img in images:
        # Make sure paths exist
        abs_path = img
        if os.path.exists(abs_path):
            img_paths.append(abs_path)
        else:
            img_paths.append(None)  # placeholder
    
    return human_prompt, img_paths, chosen_text, rejected_text

# Build Gradio app
with gr.Blocks() as demo:
    with gr.Row():
        json_dropdown = gr.Dropdown(choices=list_json_files(), label="Select JSON")
        randomize_btn = gr.Button("Randomize Entry")
    
    with gr.Row():
        prompt_text = gr.Textbox(label="Prompt", interactive=False)
    
    with gr.Row():
        img_gallery = gr.Gallery(label="Images", elem_id="image_gallery").style(grid=[3], height="auto")
    
    with gr.Row():
        chosen_textbox = gr.Textbox(label="GPT Chosen", interactive=False)
        rejected_textbox = gr.Textbox(label="GPT Rejected", interactive=False)
    
    # Connect button
    randomize_btn.click(
        fn=show_random_entry, 
        inputs=[json_dropdown], 
        outputs=[prompt_text, img_gallery, chosen_textbox, rejected_textbox]
    )

demo.launch()