Update

Dockerfile

@@ -27,7 +27,7 @@ RUN chmod +x /tmp/TMalign
 RUN huggingface-cli download westlake-repl/ProTrek_650M_UniRef50 --repo-type model --local-dir /tmp/ProTrek_650M_UniRef50
 
 # Download ProTrek faiss index
-RUN huggingface-cli download westlake-repl/ProTrek-faiss-index  --repo-type dataset  --local-dir /tmp/ProTrek-faiss-index
+RUN huggingface-cli download westlake-repl/ProTrek-faiss-index  --repo-type dataset  --local-dir /tmp/ProTrek-faiss-index --include "ProTrek_650M_UniRef50/Swiss-Prot/*"
 
 # Set up a new user named "user" with user ID 1000
 RUN useradd -m -u 1000 user

run.py

@@ -4,6 +4,8 @@ if root_dir not in sys.path:
     sys.path.append(root_dir)
 
 import gradio as gr
+import asyncio
+import os
 
 from demo.modules.search import build_search_module
 from demo.modules.compute_score import build_score_computation
@@ -17,6 +19,31 @@ with gr.Blocks() as demo:
     build_TMalign()
 
 
+# Asynchronously download the required index files
+async def run_cmd(cmd):
+    process = await asyncio.create_subprocess_shell(
+        cmd,
+        stdout=asyncio.subprocess.PIPE,
+        stderr=asyncio.subprocess.PIPE
+    )
+    stdout, stderr = await process.communicate()
+    print(f'[{cmd!r} exited with {process.returncode}]')
+    if stdout:
+        print(f'[stdout]\n{stdout.decode()}')
+    if stderr:
+        print(f'[stderr]\n{stderr.decode()}')
+
+
+async def download_faiss_index():
+    await asyncio.gather(
+        run_cmd("Hello World!"),
+        run_cmd('huggingface-cli download westlake-repl/ProTrek-faiss-index  --repo-type dataset  --local-dir /tmp/ProTrek-faiss-index --include "ProTrek_650M_UniRef50/UniRef50/*"'),
+    )
+
+
 if __name__ == '__main__':
+    # Run the async function to download the index files
+    asyncio.run(download_faiss_index())
+    
     # Run the demo
     demo.launch()