import gradio as gr
import os
from huggingface_hub import HfApi, login
from transformers import AutoTokenizer, AutoModelForCausalLM

# NVidia A10G Large (sleep after 1 hour)

# google/gemma-2-9b-it
# meta-llama/Meta-Llama-3-8B-Instruct

profile = "bstraehle"

def download_tokenizer_and_model(model_id)
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(model_id)
    model.save_pretrained(model_id)
    return tokenizer, model

def upload_model(tokenizer, model)
    model_name = model_id[model_id.rfind('/')+1:]
    print(model_name)
    model_repo_name = f"{profile}/{model_name}"
    login(token=os.environ["HF_TOKEN"])
    api = HfApi()
    api.create_repo(repo_id=model_repo_name)
    api.upload_folder(
        folder_path=model_path,
        repo_id=model_repo_name
    )
    tokenizer.push_to_hub(model_repo_name)
    
def process(model_id, dataset):
    tokenizer, model = download_model(model_id)
    upload_model(tokenizer, model)
    
    return "Done"

demo = gr.Interface(fn=process, 
                    inputs=[gr.Textbox(label = "Model ID", value = "meta-llama/Meta-Llama-3-8B-Instruct", lines = 1),
                            gr.Textbox(label = "Dataset", value = "imdb", lines = 1)],
                    outputs=[gr.Textbox(label = "Completion")])
demo.launch()