Update app.py

app.py

@@ -13,13 +13,13 @@ def load_and_preprocess_data():
 
     # Preprocess the dataset: tokenize inputs
     def preprocess_function(examples):
-        max_length = 500  # Half of 1000 for each segment
+        max_length = 4  # Now using 4 for both chosen and rejected
         inputs = {
             'input_ids': [],
             'attention_mask': [],
             'labels': []
         }
-        for example in examples['prompt']:
+        for example in examples['prompt']:  # Adjust if 'prompt' is not the key for your text
             chosen_tokenized = tokenizer(
                 example, 
                 truncation=True, 
@@ -27,7 +27,7 @@ def load_and_preprocess_data():
                 max_length=max_length
             )
             rejected_tokenized = tokenizer(
-                example, 
+                example,  
                 truncation=True, 
                 padding='max_length', 
                 max_length=max_length
@@ -36,23 +36,16 @@ def load_and_preprocess_data():
             combined_input_ids = chosen_tokenized['input_ids'] + rejected_tokenized['input_ids']
             combined_attention_mask = chosen_tokenized['attention_mask'] + rejected_tokenized['attention_mask']
     
-            # Ensure length is 1000
-            if len(combined_input_ids) > 1000:
-                combined_input_ids = combined_input_ids[:1000]
-                combined_attention_mask = combined_attention_mask[:1000]
-            else:
-                padding_length = 1000 - len(combined_input_ids)
-                combined_input_ids += [tokenizer.pad_token_id] * padding_length
-                combined_attention_mask += [0] * padding_length
-    
+            # No need for padding or truncating since we're already at max length
             inputs['input_ids'].append(combined_input_ids)
             inputs['attention_mask'].append(combined_attention_mask)
-            inputs['labels'].append(torch.tensor([1] * 500 + [0] * 500))  # Convert to tensor immediately
+            # Adjust labels to match new length; here assuming binary classification where first half is 'chosen'
+            inputs['labels'].append(torch.tensor([1] * 2 + [0] * 2))  # 4 labels for 4 tokens
     
         # Convert lists to tensors for each batch
         inputs['input_ids'] = torch.tensor(inputs['input_ids'])
         inputs['attention_mask'] = torch.tensor(inputs['attention_mask'])
-        inputs['labels'] = torch.stack(inputs['labels'])  # Stack along a new dimension
+        inputs['labels'] = torch.stack(inputs['labels'])  # Stack labels along batch dimension
     
         return inputs
     # Apply preprocessing to the dataset
@@ -105,7 +98,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=1000)  # Adjust max_length to match preprocessing
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=4)  # Adjust max_length to match preprocessing
     # Perform inference
     outputs = model(**inputs)
     prediction = torch.argmax(outputs.logits, dim=-1).item()  # Get the class with the highest score