Spaces:
Sleeping
Sleeping
| from os import makedirs, remove | |
| from os.path import exists, dirname | |
| from functools import cache | |
| import json | |
| import streamlit as st | |
| from googleapiclient.discovery import build | |
| from slugify import slugify | |
| from transformers import pipeline | |
| import uuid | |
| from beautiful_soup.beautiful_soup import get_url_content | |
| """ | |
| Request Google Search API with query and return results. | |
| """ | |
| def google_search_api_request( query ): | |
| api_key = st.secrets["google_search_api_key"] | |
| cx = st.secrets["google_search_engine_id"] | |
| service = build( | |
| "customsearch", | |
| "v1", | |
| developerKey=api_key, | |
| cache_discovery=False | |
| ) | |
| # Exclude PDFs from search results. | |
| query = query + ' -filetype:pdf' | |
| return service.cse().list( | |
| q=query, | |
| cx=cx, | |
| num=5, | |
| lr='lang_en', # lang_de | |
| fields='items(title,link),searchInformation(totalResults)' | |
| ).execute() | |
| """ | |
| Request Google Search API with query and return results. Results are cached in files. | |
| """ | |
| def search_results( query ): | |
| file_path = 'search-results/' + slugify( query ) + '.json' | |
| results = [] | |
| makedirs(dirname(file_path), exist_ok=True) | |
| if exists( file_path ): | |
| with open( file_path, 'r' ) as results_file: | |
| results = json.load( results_file ) | |
| else: | |
| search_result = google_search_api_request( query ) | |
| if int( search_result['searchInformation']['totalResults'] ) > 0: | |
| results = search_result['items'] | |
| with open( file_path, 'w' ) as results_file: | |
| json.dump( results, results_file ) | |
| if len( results ) == 0: | |
| raise Exception('No results found.') | |
| return results | |
| """ | |
| Generate summary for content. | |
| """ | |
| def generate_summary( url_id, content ): | |
| file_path = 'summaries/' + url_id + '.json' | |
| makedirs(dirname(file_path), exist_ok=True) | |
| if exists( file_path ): | |
| with open( file_path, 'r' ) as file: | |
| summary = json.load( file ) | |
| else: | |
| try: | |
| summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6") | |
| # https://huggingface.co/docs/transformers/v4.18.0/en/main_classes/pipelines#transformers.SummarizationPipeline | |
| summary = summarizer(content, max_length=130, min_length=30, do_sample=False, truncation=True) | |
| except Exception as exception: | |
| raise exception | |
| with open( file_path, 'w' ) as file: | |
| json.dump( summary, file ) | |
| return summary | |
| """ | |
| Helper function for exception notices. | |
| """ | |
| def exception_notice( exception ): | |
| query_params = st.experimental_get_query_params() | |
| if 'debug' in query_params.keys() and query_params['debug'][0] == 'true': | |
| st.exception(exception) | |
| else: | |
| st.warning(str(exception)) | |
| """ | |
| Checks if string contains keyword. | |
| """ | |
| def is_keyword_in_string( keywords, string ): | |
| for keyword in keywords: | |
| if keyword in string: | |
| return True | |
| return False | |
| def filter_strings_by_keywords( strings, keywords ): | |
| content = '' | |
| for string in strings: | |
| # Filter strings with keywords | |
| if is_keyword_in_string( keywords, string ): | |
| content += string + '\n' | |
| return content | |
| def main(): | |
| st.title('Racoon Search') | |
| query = st.text_input('Search query') | |
| query_params = st.experimental_get_query_params() | |
| if query : | |
| with st.spinner('Loading search results...'): | |
| try: | |
| results = search_results( query ) | |
| except Exception as exception: | |
| exception_notice(exception) | |
| return | |
| number_of_results = len( results ) | |
| st.success( 'Found {} results for "{}".'.format( number_of_results, query ) ) | |
| if 'debug' in query_params.keys() and query_params['debug'][0] == 'true': | |
| with st.expander("Search results JSON"): | |
| if st.button('Delete search result cache', key=query + 'cache'): | |
| remove( 'search-results/' + slugify( query ) + '.json' ) | |
| st.json( results ) | |
| progress_bar = st.progress(0) | |
| st.header('Search results') | |
| st.markdown('---') | |
| # for result in results: | |
| for index, result in enumerate(results): | |
| with st.container(): | |
| st.markdown('### ' + result['title']) | |
| url_id = uuid.uuid5( uuid.NAMESPACE_URL, result['link'] ).hex | |
| try: | |
| strings = get_url_content( result['link'] ) | |
| keywords = query.split(' ') | |
| content = filter_strings_by_keywords( strings, keywords ) | |
| # print(content) | |
| # print(len(content.split())) | |
| summary = generate_summary( url_id, content ) | |
| for sentence in summary: | |
| st.write(sentence['summary_text']) | |
| except Exception as exception: | |
| exception_notice(exception) | |
| progress_bar.progress( ( index + 1 ) / number_of_results ) | |
| col1, col2, col3 = st.columns(3) | |
| with col1: | |
| st.markdown('[Website Link]({})'.format(result['link'])) | |
| with col2: | |
| if st.button('Delete content from cache', key=url_id + 'content'): | |
| remove( 'page-content/' + url_id + '.txt' ) | |
| with col3: | |
| if st.button('Delete summary from cache', key=url_id + 'summary'): | |
| remove( 'summaries/' + url_id + '.json' ) | |
| st.markdown('---') | |
| if __name__ == '__main__': | |
| main() | |