Tag character selection and processing with DuckDB

app.py

@@ -1,7 +1,10 @@
 import gradio as gr
 import pandas as pd
+import duckdb
 
-PARQUET_FILE = 'data.parquet'
+PARQUET_FILES = {
+    "Animagine XL 3.1": 'data.parquet', 
+}
 
 preset_values = [
     "1girl, kamado nezuko, kimetsu no yaiba",
@@ -10,16 +13,27 @@ preset_values = [
     "1girl, elaina (majo no tabitabi), majo no tabitabi"
 ]
 
-def search_parquet(search_term):
+def search_data(search_term, selected_file, use_duckdb=False):
     try:
-        df = pd.read_parquet(PARQUET_FILE)
-        if search_term.strip():
-            search_results = df[df['teks'].str.lower().str.contains(search_term.lower(), na=False)]
+        if use_duckdb:
+            con = duckdb.connect(database=':memory:', read_only=False)
+            con.execute(f"CREATE TABLE my_table AS SELECT * FROM read_parquet('{PARQUET_FILES[selected_file]}')")
+            if search_term.strip():
+                results = con.execute(f"SELECT * FROM my_table WHERE lower(teks) LIKE '%{search_term.lower()}%'").fetchdf()
+            else:
+                search_terms_sql = " OR ".join([f"lower(teks) LIKE '%{term.lower()}%'" for term in preset_values])
+                results = con.execute(f"SELECT * FROM my_table WHERE {search_terms_sql}").fetchdf()
+            con.close()
         else:
-            search_results = df[df['teks'].str.lower().str.contains("|".join(preset_values).lower(), na=False)]
-        if len(search_results.columns) > 12:
-            search_results = search_results.iloc[:, :12]
-        return search_results
+            df = pd.read_parquet(PARQUET_FILES[selected_file])
+            if search_term.strip():
+                results = df[df['teks'].str.lower().str.contains(search_term.lower(), na=False)]
+            else:
+                results = df[df['teks'].str.lower().str.contains("|".join(preset_values).lower(), na=False)]
+
+        if len(results.columns) > 12:
+            results = results.iloc[:, :12]
+        return results
     except FileNotFoundError:
         return pd.DataFrame({'Error': ['Parquet file not found. Please check the file path.']})
     except Exception as e:
@@ -27,7 +41,14 @@ def search_parquet(search_term):
 
 if __name__ == "__main__":
     with gr.Blocks() as app:
-        gr.Markdown("## 🔍 Text Search for Animagine tag characters")
+        gr.Markdown("##  Text Search for Animagine tag characters")
+        with gr.Row(): #Menambahkan row untuk dropdown dan checkbox
+          file_dropdown = gr.Dropdown(
+              choices=list(PARQUET_FILES.keys()),
+              label="Select Parquet File",
+              value=list(PARQUET_FILES.keys())[0], # Default value
+          )
+          use_duckdb_checkbox = gr.Checkbox(label="Use DuckDB", value=True) # Checkbox untuk DuckDB, default true
         with gr.Column():
             search_input = gr.Textbox(
                 label="Search for characters or series:",
@@ -39,10 +60,22 @@ if __name__ == "__main__":
                 headers="auto",
             )
 
+        inputs = [search_input, file_dropdown, use_duckdb_checkbox] # Menambahkan input dropdown dan checkbox
         search_input.change(
-            fn=search_parquet,
-            inputs=search_input,
+            fn=search_data,
+            inputs=inputs,
+            outputs=search_output,
+        )
+        file_dropdown.change( #Menambahkan event change untuk dropdown
+            fn=search_data,
+            inputs=inputs,
             outputs=search_output,
         )
+        use_duckdb_checkbox.change( #Menambahkan event change untuk checkbox
+            fn=search_data,
+            inputs=inputs,
+            outputs=search_output,
+        )
+
 
     app.launch()