# Install necessary libraries
# pip install transformers datasets torch

from transformers import GPT2LMHeadModel, GPT2Tokenizer
from transformers import Trainer, TrainingArguments
from datasets import Dataset

# Step 1: Load the pre-trained GPT-2 model and tokenizer
model_name = "gpt2"  # You can use any GPT model, GPT-3, or other variants if you want a bigger model
model = GPT2LMHeadModel.from_pretrained(model_name)
tokenizer = GPT2Tokenizer.from_pretrained(model_name)

# Set padding token as GPT-2 doesn't have one by default
tokenizer.pad_token = tokenizer.eos_token

# Step 2: Prepare your training data (Instagram algorithm and feature usage)
training_data = [
    {
        "input": "How can I improve engagement on Instagram?",
        "output": "Engagement can be improved by posting at optimal times, using 20-30 relevant hashtags, and responding to comments quickly. Consider using reels for higher visibility."
    },
    {
        "input": "What are the best times to post on Instagram?",
        "output": "The best times to post on Instagram depend on your audience's time zone. Typically, posting during peak activity times such as early morning or late evening can lead to better engagement."
    },
    {
        "input": "How do I use Instagram Insights?",
        "output": "Go to your profile, tap the menu, and select 'Insights.' You can view metrics like reach, impressions, and engagement."
    },
    {
        "input": "What is the best way to use hashtags on Instagram?",
        "output": "Use a mix of trending, niche, and brand-specific hashtags. Aim for around 20-30 relevant hashtags per post. Research the most effective ones for your target audience."
    },
    {
        "input": "How can I use Instagram Stories to grow my account?",
        "output": "Instagram Stories can be used to engage your followers by sharing behind-the-scenes content, polls, Q&As, and other interactive elements. Consistency and engaging content are key."
    },
]

# Step 3: Process the data into a format suitable for training
def process_data(examples):
    # Concatenate input and output to form the training sequence
    return tokenizer(examples['input'] + tokenizer.eos_token + examples['output'], truncation=True, padding="max_length", max_length=128)

# Convert the training data into a dataset
dataset = Dataset.from_dict(training_data)
dataset = dataset.map(process_data, batched=True)

# Step 4: Split the dataset into training and validation sets
train_dataset = dataset.train_test_split(test_size=0.1)["train"]
val_dataset = dataset.train_test_split(test_size=0.1)["test"]

# Step 5: Define the training arguments
training_args = TrainingArguments(
    output_dir="./gpt2-instagram-model",     # Directory to save the model
    evaluation_strategy="epoch",              # Evaluate at the end of each epoch
    learning_rate=5e-5,                      # Learning rate for fine-tuning
    per_device_train_batch_size=4,           # Batch size for training
    per_device_eval_batch_size=4,            # Batch size for evaluation
    num_train_epochs=3,                      # Number of training epochs
    weight_decay=0.01,                       # Weight decay for regularization
    logging_dir='./logs',                    # Log directory
    logging_steps=200,                       # Log every 200 steps
)

# Step 6: Initialize the Trainer
trainer = Trainer(
    model=model,                             # The model we are training
    args=training_args,                      # Training arguments
    train_dataset=train_dataset,             # Training dataset
    eval_dataset=val_dataset,                # Validation dataset
)

# Step 7: Train the model
trainer.train()

# Step 8: Evaluate the model after training
results = trainer.evaluate()
print("Evaluation Results:", results)

# Step 9: Save the model and tokenizer
model.save_pretrained("./gpt2-instagram-model")
tokenizer.save_pretrained("./gpt2-instagram-model")

# Step 10: Use the trained model to generate responses
def generate_response(input_text):
    # Encode the input text and generate a response
    inputs = tokenizer.encode(input_text, return_tensors="pt")
    output = model.generate(inputs, max_length=100, num_return_sequences=1, no_repeat_ngram_size=2)
    
    # Decode and return the response
    response = tokenizer.decode(output[0], skip_special_tokens=True)
    return response

# Example: Generate a response
input_text = "How can I improve engagement on Instagram?"
response = generate_response(input_text)
print("Generated Response:", response)