import gradio as gr
from transformers import AutoModel
from safetensors.torch import load_file, save_file
import torch
import os

# Function to download and modify model
def add_fuzz_to_model(model_name, min_val, max_val):
    # Download model weights from Hugging Face
    model = AutoModel.from_pretrained(model_name, use_safetensors=True)
    model_file = f"{model_name}.safetensors"
    model.save_pretrained(".", safe_serialization=True)  # Save locally for access

    # Load safetensors
    weights = load_file(model_file)
    
    # Apply fuzz to each weight tensor
    def add_fuzz(tensor, min_val, max_val):
        noise = torch.empty_like(tensor).uniform_(min_val, max_val)
        return tensor + noise

    modified_weights = {}
    for key, tensor in weights.items():
        modified_weights[key] = add_fuzz(tensor, min_val, max_val)
    
    # Save modified weights
    modified_file = f"{model_name}_fuzzed.safetensors"
    save_file(modified_weights, modified_file)

    return f"Model processed and saved as {modified_file}"

# Gradio UI
def fuzz_weights_ui():
    with gr.Blocks() as interface:
        gr.Markdown("### Add Fuzz to Hugging Face Model Weights")

        with gr.Row():
            model_name = gr.Textbox(label="Hugging Face Model Name", placeholder="e.g., gpt2")
        
        with gr.Row():
            min_val = gr.Slider(-0.5, 0.0, value=-0.1, step=0.01, label="Minimum Fuzz")
            max_val = gr.Slider(0.0, 0.5, value=0.1, step=0.01, label="Maximum Fuzz")
        
        with gr.Row():
            process_button = gr.Button("Apply Fuzz")
            result = gr.Textbox(label="Result")

        # Define interaction
        process_button.click(
            fn=add_fuzz_to_model,
            inputs=[model_name, min_val, max_val],
            outputs=[result],
        )

    return interface

# Launch the app
interface = fuzz_weights_ui()
interface.launch()