split data

app.py

@@ -1,6 +1,6 @@
 import os
 import torch
-from datasets import load_dataset
+from datasets import load_dataset, DatasetDict
 from transformers import AutoTokenizer, AutoModel
 import chromadb
 import gradio as gr
@@ -15,6 +15,13 @@ def meanpooling(output, mask):
 # Load the dataset
 dataset = load_dataset("thankrandomness/mimic-iii-sample")
 
+# Split the dataset into train and validation sets
+split_dataset = dataset['train'].train_test_split(test_size=0.2, seed=42)
+dataset = DatasetDict({
+    'train': split_dataset['train'],
+    'validation': split_dataset['test']
+})
+
 # Load the model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("neuml/pubmedbert-base-embeddings-matryoshka")
 model = AutoModel.from_pretrained("neuml/pubmedbert-base-embeddings-matryoshka")

requirements.txt

@@ -1,5 +1,6 @@
 torch
 transformers
+datasets
 chromadb
 gradio
 numpy