Update app.py

app.py

@@ -54,18 +54,13 @@ client2.create_collection(
     ),
 )
 
-
-
-
-
-
 # Create function to generate embeddings (in batches) for a given dataset split
-def generate_embeddings(dataset, batch_size=32):
+def generate_embeddings(dataset, text_field, batch_size=32):
     embeddings = []
 
     with tqdm(total=len(dataset), desc=f"Generating embeddings for dataset") as pbar:
         for i in range(0, len(dataset), batch_size):
-            batch_sentences = dataset['content'][i:i+batch_size]
+            batch_sentences = dataset[text_field][i:i+batch_size]
             batch_embeddings = model.encode(batch_sentences)
             embeddings.extend(batch_embeddings)
             pbar.update(len(batch_sentences))
@@ -73,7 +68,7 @@ def generate_embeddings(dataset, batch_size=32):
     return embeddings
     
 @app.post("/uploadfile/")
-async def create_upload_file(file: UploadFile = File(...)):
+async def create_upload_file(file: UploadFile = File(...), text_field: str):
     file_savePath =  join(temp_path,file.filename)
 
     with open(file_savePath,'wb') as f:
@@ -95,7 +90,7 @@ async def create_upload_file(file: UploadFile = File(...)):
     else:
         raise NotImplementedError("This feature is not supported yet")
     # Generate and append embeddings to the train split
-    law_embeddings = generate_embeddings(full_dataset)
+    law_embeddings = generate_embeddings(full_dataset, text_field)
     full_dataset= full_dataset.add_column("embeddings", law_embeddings)
     
     if not 'uuid' in full_dataset.column_names: