|
import gradio as gr |
|
import pandas as pd |
|
import duckdb |
|
|
|
PARQUET_FILES = { |
|
"Animagine XL 3.1": 'data.parquet', |
|
} |
|
|
|
preset_values = [ |
|
"1girl, kamado nezuko, kimetsu no yaiba", |
|
"1girl, kanroji mitsuri, kimetsu no yaiba", |
|
"1girl, fern (sousou no frieren), sousou no frieren", |
|
"1girl, elaina (majo no tabitabi), majo no tabitabi" |
|
] |
|
|
|
def search_data(search_term, selected_file): |
|
try: |
|
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() |
|
|
|
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: |
|
return pd.DataFrame({'Error': [f'An error occurred: {e}']}) |
|
|
|
if __name__ == "__main__": |
|
with gr.Blocks() as app: |
|
gr.Markdown("## Text Search for Animagine tag characters") |
|
with gr.Row(): |
|
with gr.Column(): |
|
search_input = gr.Textbox( |
|
label="Search for characters or series:", |
|
placeholder="sousou no frieren", |
|
) |
|
with gr.Column(): |
|
file_dropdown = gr.Dropdown( |
|
choices=list(PARQUET_FILES.keys()), |
|
label="Select Character Tag", |
|
value=list(PARQUET_FILES.keys())[0], |
|
) |
|
search_output = gr.Dataframe( |
|
label="Search Results", |
|
value=pd.DataFrame({'Characters tag': preset_values}), |
|
headers="auto", |
|
) |
|
|
|
inputs = [search_input, file_dropdown] |
|
search_input.change( |
|
fn=search_data, |
|
inputs=inputs, |
|
outputs=search_output, |
|
) |
|
file_dropdown.change( |
|
fn=search_data, |
|
inputs=inputs, |
|
outputs=search_output, |
|
) |
|
|
|
app.launch() |