Update app.py

app.py

@@ -123,6 +123,7 @@ dataset = dataset.map(
     prepare_dataset, remove_columns=dataset.column_names,
 )
 
+dataset = dataset.train_test_split(test_size=0.1)
 
 def predict(text, speaker):
     if len(text.strip()) == 0:
@@ -135,7 +136,7 @@ def predict(text, speaker):
     input_ids = input_ids[..., :model.config.max_text_positions]
 
     ### ### ### 
-    example = dataset.iloc[0]
+    example = dataset['test'][11]
     speaker_embeddings = torch.tensor(example["speaker_embeddings"]).unsqueeze(0)
 
     speaker_embedding = torch.tensor(speaker_embedding).unsqueeze(0)