Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 from datasets import load_dataset
-from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
+from transformers import LongformerTokenizer, LongformerForSequenceClassification, Trainer, TrainingArguments
 import torch
 
 # Load and preprocess dataset
@@ -9,22 +9,22 @@ def load_and_preprocess_data():
     dataset = load_dataset('json', data_files='data.json')['train']  # Assuming single split
 
     # Load the tokenizer
-    tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
+    tokenizer = LongformerTokenizer.from_pretrained('allenai/longformer-base-4096')
 
     # Preprocess the dataset: tokenize inputs
     def preprocess_function(example):
-        # Tokenize both chosen and rejected responses
+        # Tokenize both chosen and rejected responses with a larger max_length
         chosen_tokenized = tokenizer(
             example['prompt'], 
             truncation=True, 
             padding='max_length', 
-            max_length=500
+            max_length=2000  # Adjust based on your longest example but within model limits
         )
         rejected_tokenized = tokenizer(
             example['prompt'], 
             truncation=True, 
             padding='max_length', 
-            max_length=500
+            max_length=2000  # Adjust based on your longest example but within model limits
         )
         
         # Combine chosen and rejected in one sequence
@@ -35,7 +35,7 @@ def load_and_preprocess_data():
         }
         
         # Ensure the lengths match; adjust if necessary
-        max_length = 500 * 2  # Since we're combining two sequences
+        max_length = 4096  # Longformer's max length
         if len(inputs['input_ids']) > max_length:
             inputs['input_ids'] = inputs['input_ids'][:max_length]
             inputs['attention_mask'] = inputs['attention_mask'][:max_length]
@@ -53,18 +53,18 @@ def load_and_preprocess_data():
 
 # Define and train the model
 def train_model(dataset):
-    # Load the model (BERT-based for sequence classification)
-    model = AutoModelForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)
+    # Load the Longformer model for sequence classification
+    model = LongformerForSequenceClassification.from_pretrained('allenai/longformer-base-4096', num_labels=2)
 
     # Define training arguments
     training_args = TrainingArguments(
         output_dir='./results',          # Output directory for model checkpoints
-        evaluation_strategy="epoch",    # Evaluation at the end of each epoch
-        learning_rate=2e-5,             # Learning rate
-        per_device_train_batch_size=8,  # Batch size for training
-        per_device_eval_batch_size=8,   # Batch size for evaluation
-        num_train_epochs=3,             # Number of training epochs
-        weight_decay=0.01,              # Strength of weight decay
+        evaluation_strategy="epoch",     # Evaluation at the end of each epoch
+        learning_rate=2e-5,              # Learning rate
+        per_device_train_batch_size=4,   # Reduced batch size due to longer sequences
+        per_device_eval_batch_size=4,    # Reduced batch size for evaluation
+        num_train_epochs=3,              # Number of training epochs
+        weight_decay=0.01,               # Strength of weight decay
     )
 
     # Initialize the Trainer
@@ -83,7 +83,7 @@ def train_model(dataset):
 # Gradio Interface function to predict
 def predict(text):
     # Tokenize the input text
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=4096)
     # Perform inference
     outputs = model(**inputs)
     prediction = torch.argmax(outputs.logits, dim=-1).item()  # Get the class with the highest score