Update app.py

app.py

@@ -1,7 +1,103 @@
+# 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():
@@ -47,10 +143,13 @@ def train_model(dataset):
         output_dir='./results',
         evaluation_strategy="epoch",
         learning_rate=2e-5,
-        per_device_train_batch_size=4,
-        per_device_eval_batch_size=4,
+        per_device_train_batch_size=16,  # Increased batch size for faster training
+        per_device_eval_batch_size=16,   # Increased eval batch size
         num_train_epochs=3,
         weight_decay=0.01,
+        fp16=True,  # Enable mixed precision
+        dataloader_num_workers=4,  # Increase to load data in parallel
+        gradient_accumulation_steps=2,  # Simulate larger batch size
     )
 
     class CustomTrainer(Trainer):
@@ -83,3 +182,4 @@ if __name__ == "__main__":
     dataset, tokenizer = load_and_preprocess_data()
     model = train_model(dataset)
     gr.Interface(fn=predict, inputs=gr.Textbox(), outputs="text").launch()
+