Update app.py

app.py

@@ -152,9 +152,10 @@ def prepare_training_data(qa_pairs, tokenizer):
     context = "your_bhagavad_gita_text_here"  # Replace with your preprocessed Bhagavad Gita text
     context_encoded = tokenizer(context, add_special_tokens=True, return_tensors="pt")
     # start_positions = answer_encoded.input_ids == tokenizer.convert_tokens_to_ids(tokenizer.sep_token)[0]  # Find first SEP token
-    start_positions = answer_encoded.input_ids == [tokenizer.convert_tokens_to_ids(tokenizer.sep_token)[0]]
+    # start_positions = answer_encoded.input_ids == [tokenizer.convert_tokens_to_ids(tokenizer.sep_token)[0]]
+    start_positions = answer_encoded.input_ids == [[tokenizer.convert_tokens_to_ids(tokenizer.sep_token)]]  # Double square brackets for list of list
 
-    end_positions = answer_encoded.input_ids == [tokenizer.convert_tokens_to_ids(tokenizer.eos_token)[0]]  # Find first EOS token
+    end_positions = answer_encoded.input_ids == [[tokenizer.convert_tokens_to_ids(tokenizer.eos_token)]]  # Find first EOS token
 
     # Combine all data into a dictionary for each QA pair
     encoded_data.append({