import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
import gradio as gr
from datasets import load_dataset

# Load the GPT-J model and tokenizer from Hugging Face
model_name = "EleutherAI/gpt-j-6B"  # GPT-J model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Define a function for generating responses
def generate_response(prompt):
    # Tokenize the input text
    inputs = tokenizer(prompt, return_tensors="pt")
    
    # Generate response using the model
    outputs = model.generate(
        inputs.input_ids, 
        max_length=200,      # Adjust response length
        do_sample=True,      # Sample from the model
        temperature=0.7,     # Adjust creativity (lower = more conservative)
        top_p=0.9,           # Nucleus sampling for diverse answers
        pad_token_id=tokenizer.eos_token_id  # EOS token to avoid error
    )
    
    # Decode and return the response
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Create a Gradio interface for the chatbot
def chatbot_interface(user_input):
    response = generate_response(user_input)
    return response

# Fine-tuning function
def fine_tune_model(dataset_name, output_dir='./fine-tuned-model'):
    # Load dataset (assuming it's in Hugging Face format or a similar format)
    dataset = load_dataset(dataset_name)

    # Define training arguments
    training_args = TrainingArguments(
        output_dir=output_dir,        # Output directory for the fine-tuned model
        per_device_train_batch_size=2, # Adjust batch size according to GPU capacity
        num_train_epochs=1,           # Adjust the number of epochs based on the dataset
        save_steps=500,               # Save checkpoint every 500 steps
        save_total_limit=2,           # Keep only last two checkpoints
    )

    # Define the Trainer object
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset['train'],
        eval_dataset=dataset['test'],
    )

    # Fine-tune the model
    trainer.train()

    return f"Model fine-tuned and saved to {output_dir}"

# Gradio interface for fine-tuning the model
def fine_tune_interface(dataset_name):
    result = fine_tune_model(dataset_name)
    return result

# Gradio UI for both chatbot and fine-tuning functionalities
with gr.Blocks() as demo:
    gr.Markdown("# Cable Industry Customer Service Chatbot")

    # Chatbot section
    with gr.Row():
        with gr.Column():
            chatbot_input = gr.Textbox(label="Ask the Cable Industry Chatbot")
            chatbot_output = gr.Textbox(label="Chatbot Response")
            chatbot_button = gr.Button("Generate Response")
        
        chatbot_button.click(fn=chatbot_interface, inputs=chatbot_input, outputs=chatbot_output)

    # Fine-tuning section
    with gr.Row():
        with gr.Column():
            gr.Markdown("### Fine-tune the Chatbot")
            dataset_input = gr.Textbox(label="Enter Hugging Face Dataset Name (e.g., 'cable_data')")
            fine_tune_button = gr.Button("Fine-tune Model")
            fine_tune_output = gr.Textbox(label="Fine-tune Status")
        
        fine_tune_button.click(fn=fine_tune_interface, inputs=dataset_input, outputs=fine_tune_output)

# Launch the Gradio interface
demo.launch()