change app add shutil

app.py

@@ -21,31 +21,54 @@ nltk.download('omw-1.4')
 import gdown
 import zipfile
 import os
-
+import shutil
 # Function to download and extract folder
 def download_and_extract_gdrive(file_id, destination, extract_to):
     # Download the zip file
     gdown.download(f"https://drive.google.com/uc?id={file_id}", destination, quiet=False)
     
-    # Extract the zip file
+    # Use a temporary directory to avoid conflicts
+    temp_dir = "./temp_extract"
+    os.makedirs(temp_dir, exist_ok=True)
+    
     with zipfile.ZipFile(destination, 'r') as zip_ref:
-        zip_ref.extractall(extract_to)
-    os.remove(destination)  # Clean up the downloaded zip file
-
-# Download and extract LanceDB and fine-tuned model
-st.info("Downloading and setting up necessary data. This might take a while...")
+        zip_ref.extractall(temp_dir)
+
+    # Ensure the main directory exists
+    if not os.path.exists(extract_to):
+        os.makedirs(extract_to, exist_ok=True)
+
+    # Move extracted files into the target directory
+    for item in os.listdir(temp_dir):
+        item_path = os.path.join(temp_dir, item)
+        if os.path.isdir(item_path):
+            shutil.move(item_path, extract_to)
+        elif os.path.isfile(item_path):
+            shutil.move(item_path, os.path.join(extract_to, os.path.basename(item_path)))
+    
+    # Clean up temporary files
+    shutil.rmtree(temp_dir)
+    os.remove(destination)
 
+# Download LanceDB directory
+st.info("Downloading and setting up LanceDB directory. This may take some time...")
 download_and_extract_gdrive(
-    file_id="1Qnb8bs_NXWlhDwGoswOgsp2DiLBMbfSY",  # Replace with the actual Google Drive file ID
-    destination="lancedb_directory_main",
-    extract_to="./"
+    file_id="your_lancedb_file_id",  # Replace with actual Google Drive file ID
+    destination="lancedb_directory_main.zip",
+    extract_to="./lancedb_directory_main"
 )
 
-download_and_extract_gdrive(
-    file_id="1_9VVuN_P3zsTBYzg0lAeh4ghd9zhXS3w",  # Replace with the actual Google Drive file ID
-    destination="finetuned_all_minilm_l6_v2",
-    extract_to="./"
-)
+# Ensure LanceDB structure is correct
+expected_directories = [
+    "enhanced_papers_pretrained_1.lance",
+    "enhanced_papers_pretrained_2.lance",
+    "enhanced_papers_finetuned.lance"
+]
+for subdir in expected_directories:
+    subdir_path = os.path.join("./lancedb_directory_main", subdir)
+    if not os.path.isdir(subdir_path):
+        raise FileNotFoundError(f"Missing expected directory: {subdir_path}")
+
 # # --------------------------- Load the LanceDB Table and Models --------------------------- #
 
 # Connect to LanceDB