Spaces:
Running
Running
| import json | |
| import os | |
| import re | |
| import string | |
| import traceback | |
| from typing import List, Tuple | |
| import gradio as gr | |
| import requests | |
| from huggingface_hub import HfApi | |
| hf_api = HfApi() | |
| roots_datasets = { | |
| dset.id.split("/")[-1]: dset | |
| for dset in hf_api.list_datasets( | |
| author="bigscience-data", use_auth_token=os.environ.get("bigscience_data_token") | |
| ) | |
| } | |
| def get_docid_html(docid): | |
| data_org, dataset, docid = docid.split("/") | |
| metadata = roots_datasets[dataset] | |
| locked_color = "LightGray" | |
| open_color = "#7978FF" | |
| if metadata.private: | |
| docid_html = """ | |
| <a title="This dataset is private. See the introductory text for more information" | |
| style="color:{locked_color}; font-weight: bold; text-decoration:none" | |
| onmouseover="style='color:{locked_color}; font-weight: bold; text-decoration:underline'" | |
| onmouseout="style='color:{locked_color}; font-weight: bold; text-decoration:none'" | |
| href="https://huggingface.co/datasets/bigscience-data/{dataset}" | |
| target="_blank"> | |
| 🔒{dataset} | |
| </a> | |
| <span style="color:{open_color}; ">/{docid}</span>""".format( | |
| dataset=dataset, | |
| docid=docid, | |
| locked_color=locked_color, | |
| open_color=open_color, | |
| ) | |
| else: | |
| docid_html = """ | |
| <a title="This dataset is licensed {metadata}" | |
| style="color:{open_color}; font-weight: bold; text-decoration:none" | |
| onmouseover="style='color:{open_color}; font-weight: bold; text-decoration:underline'" | |
| onmouseout="style='color:{open_color}; font-weight: bold; text-decoration:none'" | |
| href="https://huggingface.co/datasets/bigscience-data/{dataset}" | |
| target="_blank"> | |
| {dataset} | |
| </a> | |
| <span style="color:{open_color}; ">/{docid}</span>""".format( | |
| metadata=metadata.tags[0].split(":")[-1], | |
| dataset=dataset, | |
| docid=docid, | |
| open_color=open_color, | |
| ) | |
| return docid_html | |
| PII_TAGS = {"KEY", "EMAIL", "USER", "IP_ADDRESS", "ID", "IPv4", "IPv6"} | |
| PII_PREFIX = "PI:" | |
| def process_pii(text): | |
| for tag in PII_TAGS: | |
| text = text.replace( | |
| PII_PREFIX + tag, | |
| """<b><mark style="background: Fuchsia; color: Lime;">REDACTED {}</mark></b>""".format( | |
| tag | |
| ), | |
| ) | |
| return text | |
| def extract_lang_from_docid(docid): | |
| return docid.split("_")[1] | |
| def normalize(document): | |
| def remove_articles(text): | |
| return re.sub(r"\b(a|an|the)\b", " ", text) | |
| def white_space_fix(text): | |
| return " ".join(text.split()) | |
| def remove_punc(text): | |
| exclude = set(string.punctuation) | |
| return "".join(ch for ch in text if ch not in exclude) | |
| def lower(text): | |
| return text.lower() | |
| return white_space_fix(remove_articles(remove_punc(lower(document)))) | |
| def format_result(result, highlight_terms, exact_search, datasets_filter=None): | |
| text, url, docid = result | |
| if datasets_filter is not None: | |
| datasets_filter = set(datasets_filter) | |
| dataset = docid.split("/")[1] | |
| if not dataset in datasets_filter: | |
| return "" | |
| tokens_html = "" | |
| if exact_search: | |
| query_variants = [highlight_terms] | |
| # lower | |
| query_variant = highlight_terms.lower() | |
| if query_variant not in query_variants: | |
| query_variants.append(query_variant) | |
| # upper | |
| query_variant = highlight_terms.upper() | |
| if query_variant not in query_variants: | |
| query_variants.append(query_variant) | |
| # first capital | |
| query_variant = highlight_terms.lower() | |
| query_variant = query_variant[0].upper() + query_variant[1:].lower() | |
| if query_variant not in query_variants: | |
| query_variants.append(query_variant) | |
| # camel case | |
| query_tokens = highlight_terms.split() | |
| query_variant = " ".join( | |
| [token[0].upper() + token[1:].lower() for token in query_tokens] | |
| ) | |
| if query_variant not in query_variants: | |
| query_variants.append(query_variant) | |
| for query_variant in query_variants: | |
| query_start = text.find(query_variant) | |
| if query_start >= 0: | |
| query_end = query_start + len(query_variant) | |
| tokens_html = text[0:query_start] | |
| tokens_html += "<b>{}</b>".format(text[query_start:query_end]) | |
| tokens_html += text[query_end:] | |
| break | |
| else: | |
| tokens = text.split() | |
| tokens_html = [] | |
| for token in tokens: | |
| if token in highlight_terms: | |
| tokens_html.append("<b>{}</b>".format(token)) | |
| else: | |
| tokens_html.append(token) | |
| tokens_html = " ".join(tokens_html) | |
| tokens_html = process_pii(tokens_html) | |
| url_html = ( | |
| """ | |
| <span style='font-size:12px; font-family: Arial; color:Silver; text-align: left;'> | |
| <a style='text-decoration:none; color:Silver;' | |
| onmouseover="style='text-decoration:underline; color:Silver;'" | |
| onmouseout="style='text-decoration:none; color:Silver;'" | |
| href='{url}' | |
| target="_blank"> | |
| {url} | |
| </a> | |
| </span><br> | |
| """.format( | |
| url=url | |
| ) | |
| if url is not None | |
| else "" | |
| ) | |
| docid_html = get_docid_html(docid) | |
| language = extract_lang_from_docid(docid) | |
| result_html = """{} | |
| <span style='font-size:14px; font-family: Arial; color:MediumAquaMarine'>Language: {} | </span> | |
| <span style='font-size:14px; font-family: Arial; color:#7978FF; text-align: left;'>Document ID: {} | </span> | |
| <a href="https://forms.gle/AdBLLwRApqcLkHYA8" target="_blank"> | |
| <button style="color:#ffcdf8; ">🏴☠️ Flag result 🏴☠️</button> | |
| </a><br> | |
| <span style='font-family: Arial;'>{}</span><br> | |
| <br> | |
| """.format( | |
| url_html, language, docid_html, tokens_html | |
| ) | |
| return "<p>" + result_html + "</p>" | |
| def format_result_page( | |
| language, results, highlight_terms, num_results, exact_search, datasets_filter=None | |
| ) -> gr.HTML: | |
| filtered_num_results = 0 | |
| header_html = "" | |
| if language == "detect_language" and not exact_search: | |
| header_html += """<div style='font-family: Arial; color:MediumAquaMarine; text-align: center; line-height: 3em'> | |
| Detected language: <b style='color:MediumAquaMarine'>{}</b></div>""".format( | |
| list(results.keys())[0] | |
| ) | |
| result_page_html = "" | |
| for lang, results_for_lang in results.items(): | |
| print("Processing language", lang) | |
| if len(results_for_lang) == 0: | |
| if exact_search: | |
| result_page_html += """<div style='font-family: Arial; color:Silver; text-align: left; line-height: 3em'> | |
| No results found.</div>""" | |
| else: | |
| result_page_html += """<div style='font-family: Arial; color:Silver; text-align: left; line-height: 3em'> | |
| No results for language: <b>{}</b></div>""".format( | |
| lang | |
| ) | |
| continue | |
| results_for_lang_html = "" | |
| for result in results_for_lang: | |
| result_html = format_result( | |
| result, highlight_terms, exact_search, datasets_filter | |
| ) | |
| if result_html != "": | |
| filtered_num_results += 1 | |
| results_for_lang_html += result_html | |
| if language == "all" and not exact_search: | |
| results_for_lang_html = f""" | |
| <details> | |
| <summary style='font-family: Arial; color:MediumAquaMarine; text-align: left; line-height: 3em'> | |
| Results for language: <b>{lang}</b> | |
| </summary> | |
| {results_for_lang_html} | |
| </details>""" | |
| result_page_html += results_for_lang_html | |
| if num_results is not None: | |
| header_html += """<div style='font-family: Arial; color:MediumAquaMarine; text-align: center; line-height: 3em'> | |
| Total number of matches: <b style='color:MediumAquaMarine'>{}</b></div>""".format( | |
| num_results | |
| ) | |
| return header_html + result_page_html | |
| def extract_results_from_payload(query, language, payload, exact_search): | |
| results = payload["results"] | |
| processed_results = dict() | |
| datasets = set() | |
| highlight_terms = None | |
| num_results = None | |
| if exact_search: | |
| highlight_terms = query | |
| num_results = payload["num_results"] | |
| results = {"dummy": results} | |
| else: | |
| highlight_terms = payload["highlight_terms"] | |
| for lang, results_for_lang in results.items(): | |
| processed_results[lang] = list() | |
| for result in results_for_lang: | |
| text = result["text"] | |
| url = ( | |
| result["meta"]["url"] | |
| if "meta" in result | |
| and result["meta"] is not None | |
| and "url" in result["meta"] | |
| else None | |
| ) | |
| docid = result["docid"] | |
| _, dataset, _ = docid.split("/") | |
| datasets.add(dataset) | |
| processed_results[lang].append((text, url, docid)) | |
| return processed_results, highlight_terms, num_results, list(datasets) | |
| def no_query_error_message(): | |
| return f""" | |
| <p style='font-size:18px; font-family: Arial; color:MediumVioletRed; text-align: center;'> | |
| Please provide a non-empty query. | |
| </p><br><hr><br>""" | |
| def process_error(error_type, payload): | |
| if error_type == "unsupported_lang": | |
| detected_lang = payload["err"]["meta"]["detected_lang"] | |
| return f""" | |
| <p style='font-size:18px; font-family: Arial; color:MediumVioletRed; text-align: center;'> | |
| Detected language <b>{detected_lang}</b> is not supported.<br> | |
| Please choose a language from the dropdown or type another query. | |
| </p><br><hr><br>""" | |
| def extract_error_from_payload(payload): | |
| if "err" in payload: | |
| return payload["err"]["type"] | |
| return None | |
| def request_payload(query, language, exact_search, num_results=10, received_results=0): | |
| post_data = {"query": query, "k": num_results, "received_results": received_results} | |
| if language != "detect_language": | |
| post_data["lang"] = language | |
| address = ( | |
| os.environ.get("address_exact_search") | |
| if exact_search | |
| else os.environ.get("address") | |
| ) | |
| output = requests.post( | |
| address, | |
| headers={"Content-type": "application/json"}, | |
| data=json.dumps(post_data), | |
| timeout=120, | |
| ) | |
| payload = json.loads(output.text) | |
| return payload | |
| title = ( | |
| """<p style="text-align: center; font-size:28px"> 🌸 🔎 ROOTS search tool 🔍 🌸 </p>""" | |
| ) | |
| description = """ | |
| The ROOTS corpus was developed during the [BigScience workshop](https://bigscience.huggingface.co/) for the purpose | |
| of training the Multilingual Large Language Model [BLOOM](https://huggingface.co/bigscience/bloom). The ROOTS Search | |
| Tool allows you to search through the ROOTS corpus. We serve a BM25 index for each language or group of languages | |
| included in ROOTS. We also offer exact search which is enabled if you enclose your query in double quotes. More details | |
| about the implementation and use cases is available in our [paper](https://arxiv.org/abs/2302.14035) - please cite it | |
| if you use ROOTS Search Tool in your work. For more information and instructions on how to access the full corpus | |
| consult [this form](https://forms.gle/qyYswbEL5kA23Wu99).""" | |
| if __name__ == "__main__": | |
| demo = gr.Blocks(css=".underline-on-hover:hover { text-decoration: underline; }") | |
| with demo: | |
| processed_results_state = gr.State([]) | |
| highlight_terms_state = gr.State([]) | |
| num_results_state = gr.State(0) | |
| exact_search_state = gr.State(False) | |
| received_results_state = gr.State(0) | |
| with gr.Row(): | |
| gr.Markdown(value=title) | |
| with gr.Row(): | |
| gr.Markdown(value=description) | |
| with gr.Row(): | |
| query = gr.Textbox( | |
| lines=1, | |
| max_lines=1, | |
| placeholder="Put your query in double quotes for exact search.", | |
| label="Query", | |
| ) | |
| with gr.Row(): | |
| lang = gr.Dropdown( | |
| choices=[ | |
| "ar", | |
| "ca", | |
| "code", | |
| "en", | |
| "es", | |
| "eu", | |
| "fr", | |
| "id", | |
| "indic", | |
| "nigercongo", | |
| "pt", | |
| "vi", | |
| "zh", | |
| "detect_language", | |
| "all", | |
| ], | |
| value="en", | |
| label="Language", | |
| ) | |
| k = gr.Slider( | |
| 1, | |
| 100, | |
| value=10, | |
| step=1, | |
| label="Max Results in fuzzy search or Max Results per page in exact search", | |
| ) | |
| with gr.Row(): | |
| submit_btn = gr.Button("Submit") | |
| with gr.Row(visible=False) as datasets_filter: | |
| available_datasets = gr.Dropdown( | |
| type="value", | |
| choices=[], | |
| value=[], | |
| label="Datasets Filter", | |
| multiselect=True, | |
| ) | |
| with gr.Row(): | |
| result_page_html = gr.HTML(label="Results") | |
| with gr.Row(visible=False) as pagination: | |
| next_page_btn = gr.Button("Next Page") | |
| def run_query(query, lang, k, dropdown_input, received_results): | |
| query = query.strip() | |
| exact_search = False | |
| if query.startswith('"') and query.endswith('"') and len(query) >= 2: | |
| exact_search = True | |
| query = query[1:-1] | |
| else: | |
| query = " ".join(query.split()) | |
| if query == "" or query is None: | |
| return ( | |
| [], | |
| [], | |
| 0, | |
| False, | |
| no_query_error_message(), | |
| [], | |
| ) | |
| payload = request_payload(query, lang, exact_search, k, received_results) | |
| err = extract_error_from_payload(payload) | |
| if err is not None: | |
| return ( | |
| [], | |
| [], | |
| 0, | |
| False, | |
| process_error(err, payload), | |
| [], | |
| ) | |
| ( | |
| processed_results, | |
| highlight_terms, | |
| num_results, | |
| ds, | |
| ) = extract_results_from_payload( | |
| query, | |
| lang, | |
| payload, | |
| exact_search, | |
| ) | |
| result_page = format_result_page( | |
| lang, processed_results, highlight_terms, num_results, exact_search | |
| ) | |
| return ( | |
| processed_results, | |
| highlight_terms, | |
| num_results, | |
| exact_search, | |
| result_page, | |
| ds, | |
| ) | |
| def submit(query, lang, k, dropdown_input): | |
| print("submitting", query, lang, k) | |
| ( | |
| processed_results, | |
| highlight_terms, | |
| num_results, | |
| exact_search, | |
| result_page, | |
| datasets, | |
| ) = run_query(query, lang, k, dropdown_input, 0) | |
| has_more_results = exact_search and (num_results > k) | |
| current_results = ( | |
| len(next(iter(processed_results.values()))) | |
| if len(processed_results) > 0 | |
| else 0 | |
| ) | |
| return [ | |
| processed_results, | |
| highlight_terms, | |
| num_results, | |
| exact_search, | |
| gr.update(visible=True) | |
| if current_results > 0 | |
| else gr.update(visible=False), | |
| gr.Dropdown.update(choices=datasets, value=datasets), | |
| gr.update(visible=has_more_results), | |
| current_results, | |
| result_page, | |
| ] | |
| def next_page( | |
| query, | |
| lang, | |
| k, | |
| dropdown_input, | |
| received_results, | |
| processed_results, | |
| ): | |
| ( | |
| processed_results, | |
| highlight_terms, | |
| num_results, | |
| exact_search, | |
| result_page, | |
| datasets, | |
| ) = run_query(query, lang, k, dropdown_input, received_results) | |
| current_results = sum( | |
| len(results) for results in processed_results.values() | |
| ) | |
| has_more_results = exact_search and ( | |
| received_results + current_results < num_results | |
| ) | |
| print("received_results", received_results) | |
| print("current_results", current_results) | |
| print("has_more_results", has_more_results) | |
| return [ | |
| processed_results, | |
| highlight_terms, | |
| num_results, | |
| exact_search, | |
| gr.update(visible=True) | |
| if current_results > 0 | |
| else gr.update(visible=False), | |
| gr.Dropdown.update(choices=datasets, value=datasets), | |
| gr.update(visible=current_results >= k and has_more_results), | |
| received_results + current_results, | |
| result_page, | |
| ] | |
| def filter_datasets( | |
| lang, | |
| processed_results, | |
| highlight_terms, | |
| num_results, | |
| exact_search, | |
| datasets_filter, | |
| ): | |
| result_page_html = format_result_page( | |
| lang, | |
| processed_results, | |
| highlight_terms, | |
| num_results, | |
| exact_search, | |
| datasets_filter, | |
| ) | |
| return result_page_html | |
| query.submit( | |
| fn=submit, | |
| inputs=[query, lang, k, available_datasets], | |
| outputs=[ | |
| processed_results_state, | |
| highlight_terms_state, | |
| num_results_state, | |
| exact_search_state, | |
| datasets_filter, | |
| available_datasets, | |
| pagination, | |
| received_results_state, | |
| result_page_html, | |
| ], | |
| ) | |
| submit_btn.click( | |
| submit, | |
| inputs=[query, lang, k, available_datasets], | |
| outputs=[ | |
| processed_results_state, | |
| highlight_terms_state, | |
| num_results_state, | |
| exact_search_state, | |
| datasets_filter, | |
| available_datasets, | |
| pagination, | |
| received_results_state, | |
| result_page_html, | |
| ], | |
| ) | |
| next_page_btn.click( | |
| next_page, | |
| inputs=[ | |
| query, | |
| lang, | |
| k, | |
| available_datasets, | |
| received_results_state, | |
| processed_results_state, | |
| ], | |
| outputs=[ | |
| processed_results_state, | |
| highlight_terms_state, | |
| num_results_state, | |
| exact_search_state, | |
| datasets_filter, | |
| available_datasets, | |
| pagination, | |
| received_results_state, | |
| result_page_html, | |
| ], | |
| ) | |
| available_datasets.change( | |
| filter_datasets, | |
| inputs=[ | |
| lang, | |
| processed_results_state, | |
| highlight_terms_state, | |
| num_results_state, | |
| exact_search_state, | |
| available_datasets, | |
| ], | |
| outputs=result_page_html, | |
| ) | |
| demo.launch(enable_queue=False, debug=True) | |