import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import os

@spaces.GPU(duration=120)
def convert(token, model_id):
    try:
        os.environ["HF_TOKEN"] = token
        model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32)
        config = model.config
        del config.quantization_config
        del config._pre_quantization_dtype
        model.config = config
        model = model.dequantize()
        
        tokenizer = AutoTokenizer.from_pretrained(model_id)
        
        output_dir = model_id.split("/")[-1] + "-" + "dequantized"
        model.push_to_hub(output_dir)
        tokenizer.push_to_hub(output_dir)
        
        os.environ["HF_TOKEN"] = "hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx" # security practice
        return "Successfully pushed model to hub!"
    except Exception as e:
        os.environ["HF_TOKEN"] = "hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx" # security practice part 2
        return f"An error occured: {e}"

demo = gr.Interface(fn=convert, inputs=[gr.Textbox(placeholder="hf_..."), gr.Textbox(placeholder="user/model")], outputs=gr.Text())
demo.launch()