Update app.py

app.py

@@ -1,64 +1,63 @@
 import gradio as gr
-from datasets import load_dataset, Dataset
+from datasets import load_dataset
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
 import torch
 
+# Load and preprocess dataset
 def load_and_preprocess_data():
     # Load the dataset from the uploaded JSON file
-    dataset = load_dataset('json', data_files='data.json')
-
-    # Extract "prompt" and "chosen"/"rejected" as input-output pairs
-    def process_row(example):
-        prompts = []
-        chosen_responses = []
-        rejected_responses = []
-        
-        for row in example['rows']:
-            prompts.append(row['row']['prompt'])
-            chosen_responses.append(row['row']['chosen'])
-            rejected_responses.append(row['row']['rejected'])
-        
-        return {
-            'prompt': prompts,
-            'chosen': chosen_responses,
-            'rejected': rejected_responses
-        }
-
-    # Apply the processing function to the dataset
-    dataset = dataset.map(process_row, batched=True)
+    dataset = load_dataset('json', data_files='data.json')['train']  # Assuming single split
 
+    # The dataset already has 'prompt', 'chosen', and 'rejected'. We can preprocess directly.
     # Load the tokenizer
     tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
 
     # Preprocess the dataset: tokenize inputs
-    def preprocess_function(examples):
-        return tokenizer(examples['prompt'], truncation=True, padding='max_length', max_length=512)
+    def preprocess_function(example):
+        inputs = tokenizer(
+            example['prompt'], 
+            truncation=True, 
+            padding='max_length', 
+            max_length=512
+        )
+        # Assume binary classification: 1 for chosen, 0 for rejected
+        inputs['labels'] = [1] * len(example['chosen']) + [0] * len(example['rejected'])
+        inputs['input_ids'] = (
+            inputs['input_ids'][:len(example['chosen'])] + 
+            inputs['input_ids'][len(example['chosen']):]
+        )
+        return inputs
 
     # Apply preprocessing to the dataset
     dataset = dataset.map(preprocess_function, batched=True)
 
+    # Convert dataset to PyTorch format
+    dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])
+
     return dataset, tokenizer
+
+# 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)  # Change num_labels if needed
+    model = AutoModelForSequenceClassification.from_pretrained('bert-base-uncased', 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=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
     )
 
     # Initialize the Trainer
     trainer = Trainer(
         model=model,
         args=training_args,
-        train_dataset=dataset['train'],
-        eval_dataset=dataset['test'],  # If you have a test split, otherwise just use train for both
+        train_dataset=dataset,
+        eval_dataset=dataset,  # You can split the dataset for evaluation if needed
     )
 
     # Train the model
@@ -68,10 +67,12 @@ 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)
+    # Perform inference
     outputs = model(**inputs)
     prediction = torch.argmax(outputs.logits, dim=-1).item()  # Get the class with the highest score
-    return f"Prediction: {prediction}"
+    return f"Prediction: {'chosen' if prediction == 1 else 'rejected'}"
 
 # Main script to run training
 if __name__ == "__main__":