Fix download file name, Add Delete Collection Button

app.py

@@ -32,7 +32,7 @@ async def initialize_model() -> Dict:
         os.makedirs("./pretrained/colpaligemma-3b-mix-448-base", exist_ok=True)
         files_to_download = ["adapter_model.safetensors",
                             "config.json",
-                            "model-00002-of-00002.safetensors",
+                            "model-00001-of-00002.safetensors",
                             "model-00002-of-00002.safetensors",
                             "preprocessor_config.json",
                             "tokenizer.json",
@@ -177,6 +177,16 @@ async def search_with_llm(query: str,
     
     return response.response, [Image.open(BytesIO(base64.b64decode(image))) for image in response.source_images]
 
+async def delete_collection(target_collection):
+    if await model_dict["vector_store_client"].collection_exists(collection_name=target_collection):
+        await model_dict["vector_store_client"].delete_collection(collection_name=target_collection, timeout=100)
+        choices = await get_collection_names(model_dict["vector_store_client"])
+        return (f"Deleted collection {target_collection}", gr.Dropdown(choices=choices), choices)
+    else:
+        choices = await get_collection_names(model_dict["vector_store_client"])
+        return (f"Collection {target_collection} is not found.", gr.Dropdown(choices=choices), choices)
+        
+
 
 def build_gui():
     with gr.Blocks() as demo:
@@ -199,10 +209,15 @@ def build_gui():
                 message_box = gr.Textbox(value="File not yet uploaded", 
                                         show_label=False,
                                         interactive=False)
-                convert_button = gr.Button("🔄 Convert and upload")
+                with gr.Row(equal_height=True):
+                    delete_button = gr.Button("🗑️ Delete collection")
+                    convert_button = gr.Button("🔄 Convert and upload")
                 
                 # Define the actions for conversion
                 convert_button.click(index, inputs=[files, target_collection], outputs=[message_box, target_collection, choices])
+                
+                # Define the actions for delete collection
+                delete_button.click(delete_collection, inputs=[target_collection], outputs=[message_box, target_collection, choices])
                     
             with gr.Column():
                 gr.Markdown("## 3️. Enter your question")