import gradio as gr
from transformers import GPT2Tokenizer, GPT2LMHeadModel, Trainer, TrainingArguments, DataCollatorForLanguageModeling
from datasets import Dataset
import os

# Load the GPT-2 model and tokenizer
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
model = GPT2LMHeadModel.from_pretrained("gpt2")

# Add a padding token to the tokenizer
tokenizer.pad_token = tokenizer.eos_token

# Load and format your custom dataset for Q&A training
def load_custom_dataset(file_path):
    with open(file_path, "r", encoding="utf-8") as file:
        lines = file.readlines()
    questions, answers = [], []
    for line in lines:
        if line.startswith("Q:"):
            questions.append(line.strip())
        elif line.startswith("A:"):
            answers.append(line.strip())
    print(f"Loaded {len(questions)} questions and {len(answers)} answers.")
    return Dataset.from_dict({"questions": questions, "answers": answers})


dataset_file = "stroon_tutor.txt"  # Make sure your file is formatted correctly
if os.path.exists(dataset_file):
    dataset = load_custom_dataset(dataset_file)
else:
    raise FileNotFoundError(f"File {dataset_file} not found.")

# Tokenize the dataset with distinct question and answer tokens
def tokenize_qa(examples):
    input_texts = [q for q in examples["questions"]]
    target_texts = [a for a in examples["answers"]]
    model_inputs = tokenizer(input_texts, max_length=128, truncation=True, padding="max_length")
    labels = tokenizer(target_texts, max_length=128, truncation=True, padding="max_length")
    model_inputs["labels"] = labels["input_ids"]
    return model_inputs

tokenized_dataset = dataset.map(tokenize_qa, batched=True)

# Check the size of the tokenized dataset
print(f"Tokenized dataset has {len(tokenized_dataset)} entries.")  # Print dataset size after tokenization

# Set up training arguments
training_args = TrainingArguments(
    output_dir="./results",                   
    evaluation_strategy="no",                 
    learning_rate=1e-5,                       
    per_device_train_batch_size=2,            
    num_train_epochs=3,                       
    weight_decay=0.01,                        
)

# Define the Trainer
data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    data_collator=data_collator
)

# Fine-tune the model
trainer.train()

# Function to generate responses from the fine-tuned model
def generate_response(prompt):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(
        **inputs, 
        max_length=50, 
        num_beams=5,               
        no_repeat_ngram_size=2,    
        do_sample=True,            
        top_k=50,                  
        top_p=0.95                 
    )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# Set up Gradio interface for interaction
iface = gr.Interface(fn=generate_response, inputs="text", outputs="text", title="GPT-2 Fine-Tuned Assistant")
iface.launch()