import json
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
from huggingface_hub import HfApi
import gradio as gr
import time

# Global variables to store the model and tokenizer
merged_model = None
tokenizer = None

def merge_models():
    print("Loading base model...")
    base_model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B")
    
    print("Loading fine-tuned model...")
    fine_tuned_model_path = "NoaiGPT/autotrain-14mrs-fc44l"
    fine_tuned_model = PeftModel.from_pretrained(base_model, fine_tuned_model_path)
    
    print("Merging models...")
    merged_model = fine_tuned_model.merge_and_unload()
    
    print("Saving merged model...")
    merged_model.save_pretrained("merged_model")
    
    print("Saving tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
    tokenizer.save_pretrained("merged_model")
    
    return merged_model, tokenizer

def push_to_hub(repo_name):
    print(f"Pushing merged model to Hugging Face Hub: {repo_name}")
    api = HfApi()
    
    try:
        api.create_repo(repo_name, private=True)
        print(f"Created new repository: {repo_name}")
    except Exception as e:
        print(f"Repository already exists or error occurred: {e}")
    
    api.upload_folder(
        folder_path="merged_model",
        repo_id=repo_name,
        repo_type="model",
    )
    print("Model pushed successfully!")

def generate_text(input_text):
    global merged_model, tokenizer
    
    if merged_model is None or tokenizer is None:
        return "Model not loaded. Please run the merge process first."
    
    input_ids = tokenizer.encode(input_text, return_tensors="pt")
    
    with torch.no_grad():
        output = merged_model.generate(input_ids, max_length=200, num_return_sequences=1)
    
    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
    return generated_text

def run_merge_and_push():
    global merged_model, tokenizer
    
    # Merge models
    merged_model, tokenizer = merge_models()
    
    # Generate a unique repository name
    timestamp = int(time.time())
    hub_repo_name = f"NoaiGPT/merged-llama3-8b-instruct-{timestamp}"
    
    # Push to Hugging Face Hub
    push_to_hub(hub_repo_name)
    
    return f"Model merged and pushed to Hugging Face Hub successfully! Repository: {hub_repo_name}"

# Gradio interface
with gr.Blocks() as demo:
    with gr.Tab("Generate Text"):
        iface = gr.Interface(
            fn=generate_text,
            inputs=gr.Textbox(lines=5, label="Input Text"),
            outputs=gr.Textbox(lines=10, label="Generated Text"),
            title="Merged Llama 3 8B Instruct Model",
            description="Enter text to generate a response from the merged model.",
        )
        iface.render()
    
    with gr.Tab("Merge and Push"):
        merge_and_push_button = gr.Button("Merge Models and Push to Hub")
        merge_and_push_output = gr.Textbox(lines=20, label="Merge and Push Output")
        merge_and_push_button.click(run_merge_and_push, outputs=merge_and_push_output)

# Launch the interface
if __name__ == "__main__":
    demo.launch()