import gradio as gr
import json
import os

def convert_index(file):
    with open(file.name, 'r') as f:
        data = json.load(f)
    
    new_data = {k: v for k, v in data.items()}
    new_data['weight_map'] = {k: v.replace('pytorch_model', 'model').replace('.bin', '.safetensors') for k, v in data['weight_map'].items()}
    
    output_path = os.path.join(os.path.dirname(file.name), 'model.safetensors.index.json')
    with open(output_path, 'w') as f:
        json.dump(new_data, f, indent=4)
    
    return output_path

iface = gr.Interface(
    fn=convert_index,
    inputs=gr.File(label="Upload pytorch_model.bin.index.json"),
    outputs=gr.File(label="Download model.safetensors.index.json"),
    title="Convert PyTorch Index to SafeTensors Index",
    description="Upload a pytorch_model.bin.index.json file to convert it to model.safetensors.index.json"
)

iface.launch()