# import gradio as gr
# from datasets import load_dataset
# from transformers import LongformerTokenizer, LongformerForSequenceClassification, Trainer, TrainingArguments
# import torch

# # Load and preprocess dataset
# def load_and_preprocess_data():
#     dataset = load_dataset('json', data_files='data.json')['train']

#     tokenizer = LongformerTokenizer.from_pretrained('allenai/longformer-base-4096')

#     def preprocess_function(examples):
#         inputs = {
#             'input_ids': [],
#             'attention_mask': [],
#             'labels': []
#         }
#         for example in examples['prompt']:
#             # Tokenize input
#             tokenized = tokenizer(
#                 example, 
#                 truncation=True, 
#                 padding='max_length', 
#                 max_length=128  # Increase max_length to suit model needs
#             )

#             inputs['input_ids'].append(tokenized['input_ids'])
#             inputs['attention_mask'].append(tokenized['attention_mask'])
#             inputs['labels'].append(1 if "chosen" in example else 0)  # Assign binary labels

#         inputs['input_ids'] = torch.tensor(inputs['input_ids'], dtype=torch.long)
#         inputs['attention_mask'] = torch.tensor(inputs['attention_mask'], dtype=torch.long)
#         inputs['labels'] = torch.tensor(inputs['labels'], dtype=torch.long)

#         return inputs

#     dataset = dataset.map(preprocess_function, batched=True)
#     dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])

#     return dataset, tokenizer

# # Define and train the model
# def train_model(dataset):
#     model = LongformerForSequenceClassification.from_pretrained('allenai/longformer-base-4096', num_labels=2)

#     training_args = TrainingArguments(
#         output_dir='./results',
#         evaluation_strategy="epoch",
#         learning_rate=2e-5,
#         per_device_train_batch_size=4,
#         per_device_eval_batch_size=4,
#         num_train_epochs=3,
#         weight_decay=0.01,
#     )

#     class CustomTrainer(Trainer):
#         def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
#             labels = inputs.pop("labels")
#             outputs = model(**inputs)
#             logits = outputs.logits
#             loss = torch.nn.CrossEntropyLoss()(logits.view(-1, model.config.num_labels), labels.view(-1))
#             return (loss, outputs) if return_outputs else loss

#     trainer = CustomTrainer(
#         model=model,
#         args=training_args,
#         train_dataset=dataset,
#         eval_dataset=dataset,
#     )

#     trainer.train()
#     return model

# # Gradio Interface function to predict
# def predict(text):
#     inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)
#     outputs = model(**inputs)
#     prediction = torch.argmax(outputs.logits, dim=-1).item()
#     return f"Prediction: {'chosen' if prediction == 1 else 'rejected'}"

# # Main script
# if __name__ == "__main__":
#     dataset, tokenizer = load_and_preprocess_data()
#     model = train_model(dataset)
#     gr.Interface(fn=predict, inputs=gr.Textbox(), outputs="text").launch()










import gradio as gr
from datasets import load_dataset
from transformers import LongformerTokenizer, LongformerForSequenceClassification, Trainer, TrainingArguments
import torch
from torch.utils.data import DataLoader

# Load and preprocess dataset
def load_and_preprocess_data():
    dataset = load_dataset('json', data_files='data.json')['train']

    tokenizer = LongformerTokenizer.from_pretrained('allenai/longformer-base-4096')

    def preprocess_function(examples):
        inputs = {
            'input_ids': [],
            'attention_mask': [],
            'labels': []
        }
        for example in examples['prompt']:
            # Tokenize input
            tokenized = tokenizer(
                example, 
                truncation=True, 
                padding='max_length', 
                max_length=128  # Increase max_length to suit model needs
            )

            inputs['input_ids'].append(tokenized['input_ids'])
            inputs['attention_mask'].append(tokenized['attention_mask'])
            inputs['labels'].append(1 if "chosen" in example else 0)  # Assign binary labels

        inputs['input_ids'] = torch.tensor(inputs['input_ids'], dtype=torch.long)
        inputs['attention_mask'] = torch.tensor(inputs['attention_mask'], dtype=torch.long)
        inputs['labels'] = torch.tensor(inputs['labels'], dtype=torch.long)

        return inputs

    dataset = dataset.map(preprocess_function, batched=True)
    dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])

    return dataset, tokenizer

# Define and train the model
def train_model(dataset):
    model = LongformerForSequenceClassification.from_pretrained('allenai/longformer-base-4096', num_labels=2)

    training_args = TrainingArguments(
        output_dir='./results',
        evaluation_strategy="epoch",
        learning_rate=2e-5,
        per_device_train_batch_size=2,  # Try reducing the batch size
        per_device_eval_batch_size=2,
        gradient_accumulation_steps=2,  # Accumulate gradients over 2 steps
        num_train_epochs=3,
        weight_decay=0.01,
    )


    class CustomTrainer(Trainer):
        def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
            labels = inputs.pop("labels")
            outputs = model(**inputs)
            logits = outputs.logits
            loss = torch.nn.CrossEntropyLoss()(logits.view(-1, model.config.num_labels), labels.view(-1))
            return (loss, outputs) if return_outputs else loss

    trainer = CustomTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
        eval_dataset=dataset,
    )

    trainer.train()
    return model

# Gradio Interface function to predict
def predict(text):
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)
    outputs = model(**inputs)
    prediction = torch.argmax(outputs.logits, dim=-1).item()
    return f"Prediction: {'chosen' if prediction == 1 else 'rejected'}"

# Main script
if __name__ == "__main__":
    dataset, tokenizer = load_and_preprocess_data()
    model = train_model(dataset)
    gr.Interface(fn=predict, inputs=gr.Textbox(), outputs="text").launch()