import streamlit as st import pandas as pd import os import numpy as np from sentence_transformers import SentenceTransformer, models import torch from sentence_transformers.quantization import semantic_search_faiss from pathlib import Path import time import plotly.express as px import doi import requests from datetime import datetime, timedelta API_URL = ( "https://api-inference.huggingface.co/models/mixedbread-ai/mxbai-embed-large-v1" ) from openai import OpenAI api_key = os.getenv('API_KEY') base_url = os.getenv("BASE_URL") client_openai = OpenAI( api_key=api_key, base_url=base_url, ) api_key_kimi = os.getenv('API_KEY_KIMI') base_url_kimi = os.getenv("BASE_URL_KIMI") client_openai_kimi = OpenAI( api_key=api_key_kimi, base_url=base_url_kimi, ) API_TOKEN = os.getenv('hf_token') headers = {"Authorization": f"Bearer {API_TOKEN}"} def query_hf_api(text, api=API_URL, parameters=None): if not parameters: payload = {"inputs": text} else: payload = { "inputs": text, "parameters": parameters, } response = requests.post(api, headers=headers, json=payload) try: response_data = response.json() except requests.exceptions.JSONDecodeError: st.error("Failed to get a valid response from the server. Please try again later.") return {} # Prepare an empty placeholder that can be filled if needed progress_placeholder = st.empty() # Check if the model is currently loading if "error" in response_data and "loading" in response_data["error"]: estimated_time = response_data.get("estimated_time", 30) # Default wait time to 30 seconds if not provided with progress_placeholder.container(): st.warning( f"Model from :hugging_face: is currently loading. Estimated wait time: {estimated_time:.1f} seconds. Please wait...") # Create a progress bar within the container progress_bar = st.progress(0) for i in range(int(estimated_time) + 5): # Adding a buffer time to ensure the model is loaded # Update progress bar. The factor of 100 is used to convert to percentage completion progress = int((i / (estimated_time + 5)) * 100) progress_bar.progress(progress) time.sleep(1) # Wait for a second # Clear the placeholder once loading is complete progress_placeholder.empty() st.rerun() # Rerun the app after waiting return response_data def normalize_embeddings(embeddings): """ Normalizes the embeddings matrix, so that each sentence embedding has unit length. Args: embeddings (Tensor): The embeddings tensor to normalize. Returns: Tensor: The normalized embeddings. """ if embeddings.dim() == 1: # Add an extra dimension if the tensor is 1-dimensional embeddings = embeddings.unsqueeze(0) return torch.nn.functional.normalize(embeddings, p=2, dim=1) def quantize_embeddings( embeddings, precision="ubinary", ranges=None, calibration_embeddings=None ): """ Quantizes embeddings to a specified precision using PyTorch and numpy. Args: embeddings (Tensor): The embeddings to quantize, assumed to be a Tensor. precision (str): The precision to convert to. ranges (np.ndarray, optional): Ranges for quantization. calibration_embeddings (Tensor, optional): Embeddings used for calibration. Returns: Tensor: The quantized embeddings. """ if precision == "float32": return embeddings.float() if precision in ["int8", "uint8"]: if ranges is None: if calibration_embeddings is not None: ranges = torch.stack( ( torch.min(calibration_embeddings, dim=0)[0], torch.max(calibration_embeddings, dim=0)[0], ) ) else: ranges = torch.stack( (torch.min(embeddings, dim=0)[0], torch.max(embeddings, dim=0)[0]) ) starts, ends = ranges[0], ranges[1] steps = (ends - starts) / 255 if precision == "uint8": quantized_embeddings = torch.clip( ((embeddings - starts) / steps), 0, 255 ).byte() elif precision == "int8": quantized_embeddings = torch.clip( ((embeddings - starts) / steps - 128), -128, 127 ).char() elif precision == "binary" or precision == "ubinary": embeddings_np = embeddings.numpy() > 0 packed_bits = np.packbits(embeddings_np, axis=-1) if precision == "binary": quantized_embeddings = torch.from_numpy(packed_bits - 128).char() else: quantized_embeddings = torch.from_numpy(packed_bits).byte() else: raise ValueError(f"Precision {precision} is not supported") return quantized_embeddings def process_embeddings(embeddings, precision="ubinary", calibration_embeddings=None): """ Normalizes and quantizes embeddings from an API list to a specified precision using PyTorch. Args: embeddings (list or Tensor): Raw embeddings from an external API, either as a list or a Tensor. precision (str): Desired precision for quantization. calibration_embeddings (Tensor, optional): Embeddings for calibration. Returns: Tensor: Processed embeddings, normalized and quantized. """ # Convert list to Tensor if necessary if isinstance(embeddings, list): embeddings = torch.tensor(embeddings, dtype=torch.float32) elif not isinstance(embeddings, torch.Tensor): st.error(embeddings) raise TypeError( f"Embeddings must be a list or a torch.Tensor. Message from the server: {embeddings}" ) # Convert calibration_embeddings list to Tensor if necessary if isinstance(calibration_embeddings, list): calibration_embeddings = torch.tensor( calibration_embeddings, dtype=torch.float32 ) elif calibration_embeddings is not None and not isinstance( calibration_embeddings, torch.Tensor ): raise TypeError( "Calibration embeddings must be a list or a torch.Tensor if provided. " ) normalized_embeddings = normalize_embeddings(embeddings) quantized_embeddings = quantize_embeddings( normalized_embeddings, precision=precision, calibration_embeddings=calibration_embeddings, ) return quantized_embeddings.cpu().numpy() # Load data and embeddings @st.cache_resource(ttl="1d") def load_data_embeddings(): existing_data_path = "aggregated_data" new_data_directory_bio = "db_update" existing_embeddings_path = "biorxiv_ubin_embaddings.npy" updated_embeddings_directory_bio = "embed_update" new_data_directory_med = "db_update_med" updated_embeddings_directory_med = "embed_update_med" # Load existing database and embeddings df_existing = pd.read_parquet(existing_data_path) embeddings_existing = np.load(existing_embeddings_path, allow_pickle=True) print(f"Existing data shape: {df_existing.shape}, Existing embeddings shape: {embeddings_existing.shape}") # Determine the embedding size from existing embeddings embedding_size = embeddings_existing.shape[1] # Prepare lists to collect new updates df_updates_list = [] embeddings_updates_list = [] # Helper function to process updates from a specified directory def process_updates(new_data_directory, updated_embeddings_directory): new_data_files = sorted(Path(new_data_directory).glob("*.parquet")) print(new_data_files) for data_file in new_data_files: corresponding_embedding_file = Path(updated_embeddings_directory) / ( data_file.stem + ".npy" ) if corresponding_embedding_file.exists(): df = pd.read_parquet(data_file) new_embeddings = np.load(corresponding_embedding_file, allow_pickle=True) # Check if the number of rows in the DataFrame matches the number of rows in the embeddings if df.shape[0] != new_embeddings.shape[0]: print( f"Shape mismatch for {data_file.name}: DataFrame has {df.shape[0]} rows, embeddings have {new_embeddings.shape[0]} rows. Skipping.") continue # Check embedding size and adjust if necessary if new_embeddings.shape[1] != embedding_size: print(f"Skipping {data_file.name} due to embedding size mismatch.") continue df_updates_list.append(df) embeddings_updates_list.append(new_embeddings) else: print(f"No corresponding embedding file found for {data_file.name}") # Process updates from both BioRxiv and MedArXiv process_updates(new_data_directory_bio, updated_embeddings_directory_bio) process_updates(new_data_directory_med, updated_embeddings_directory_med) # Concatenate all updates if df_updates_list: df_updates = pd.concat(df_updates_list) else: df_updates = pd.DataFrame() if embeddings_updates_list: embeddings_updates = np.vstack(embeddings_updates_list) else: embeddings_updates = np.array([]) # Append new data to existing, handling duplicates as needed df_combined = pd.concat([df_existing, df_updates]) # Create a mask for filtering mask = ~df_combined.duplicated(subset=["title"], keep="last") df_combined = df_combined[mask] # Combine embeddings, ensuring alignment with the DataFrame embeddings_combined = ( np.vstack([embeddings_existing, embeddings_updates]) if embeddings_updates.size else embeddings_existing ) # Filter the embeddings based on the dataframe unique entries embeddings_combined = embeddings_combined[mask] return df_combined, embeddings_combined LLM_prompt = "Review the abstracts listed above and create a list and summary that captures their main themes and findings. Identify any commonalities across the abstracts and highlight these in your summary. Ensure your response is concise, avoids external links, and is formatted in markdown.\n\n" def summarize_abstract(abstract, llm_model="llama-3.1-70b-versatile", instructions=LLM_prompt, api_key=""): """ Summarizes the provided abstract using a specified LLM model. Parameters: - abstract (str): The abstract text to be summarized. - llm_model (str): The LLM model used for summarization. Defaults to "llama-3.1-70b-versatile". Returns: - str: A summary of the abstract, condensed into one to two sentences. """ print("use openai api: gpt-4o-mini") client = client_openai formatted_text = "\n".join(f"{idx + 1}. {abstract}" for idx, abstract in enumerate(abstracts)) try: # Create a chat completion with the abstract and specified LLM model chat_completion = client.chat.completions.create( messages=[{"role": "user", "content": f'"{formatted_text}" {instructions}'}], model="gpt-4o-mini", ) except Exception as e: # Catch the exception print(f"An error occurred: {e}") # Print the error return 'LLM API not available or above the usage limit.' # Return the summarized content return chat_completion.choices[0].message.content def summarize_abstract_kimi(title, link): """ Summarizes the provided abstract using a specified LLM model. Parameters: - abstract (str): The abstract text to be summarized. - llm_model (str): The LLM model used for summarization. Defaults to "llama-3.1-70b-versatile". Returns: - str: A summary of the abstract, condensed into one to two sentences. """ print("use openai api: moonshot-v1-32k") print(title, link) client = client_openai_kimi formatted_text = "The paper we are going to discuss is "+ title +". The link is"+link+""" . Please use this as a basis to answer my questions. Please output your answers according to the following format. Please pay attention to the logic of subheading stratification and ensure that each layer includes 4-10 points. **Q: What problem does this paper try to solve?** A: [Use one sentence to summarize what problem this paper tries to solve] 1. Subheading 1: [Content under subheading 1] 2. Subheading 2: [Content under subheading 2] 3. Subheading 3: […] […] ** Q: What are the related studies?** A: [Use one sentence to summarize the relevant research] 1. Subheading 1: [Subheading 1] 2. Subheading 2: [Subheading 2] 3. Subheading 3: […] […] ** Q: How does the paper solve this problem?** A: [Use one sentence to summarize how the paper solves this problem] 1. Subheading 1: [Subheading 1] 2. Subheading 2: [Subheading 2] 3. Subheading 3: […] […] ** Q: What experiments were done in the paper?** A: [Use one sentence to summarize the experiments done in the paper] 1. Subheading 1: [Subheading 1] 2. Subheading 2: [Subheading 2] 3. Subheading 3: […] […] ** Q: Is there anything that can be further explored?** A: [Use one sentence here to summarize what can be further explored] 1. Subheading 1: [Content under subheading 1] 2. Subheading 2: [Content under subheading 2] 3. Subheading 3: […] […] ** Q: Summarize the main content of the paper** A: [Use one sentence here to summarize the main content of the paper] 1. Research background: […] 2. Research methods: […] 3. Experimental design: […] 4. Main findings: […] 5. Research contributions: […] 6. Future research directions: […] 7. Methods and tools: […] 8. Dataset: […] 9. Conclusion: […] """ try: # Create a chat completion with the abstract and specified LLM model chat_completion = client.chat.completions.create( messages=[{"role": "user", "content": f'"{formatted_text}"'}], model="moonshot-v1-32k", ) except Exception as e: # Catch the exception print(f"An error occurred: {e}") # Print the error return 'LLM API not available or above the usage limit.' # Return the summarized content return chat_completion.choices[0].message.content def define_style(): st.markdown( """ """, unsafe_allow_html=True, ) def logo(db_update_date, db_size_bio, db_size_med): # Initialize Streamlit app biorxiv_logo = "https://www.biorxiv.org/sites/default/files/biorxiv_logo_homepage.png" medarxiv_logo = "https://www.medrxiv.org/sites/default/files/medRxiv_homepage_logo.png" st.markdown( f"""
BioRxiv logo
medRxiv logo

LLM-based search and summary tool for bioRxiv & medRxiv

How to use:
1: Enter your search query (Optional modification "Top k results to display")
2: Press Enter in the query box (or click the search button) to search.
3: When the search results are displayed, you can click on the line of interest to view the overview information of this article.
4: If you want to learn more about the paper, you can jump to the paper pdf by clicking on 'Full Text Read' link.
5: Enter summary prompt in the prompt input box.
6: Click "AI summary" to summarize the search results above.

Last database update: {db_update_date}; Database size: bioRxiv: {db_size_bio} / medRxiv: {db_size_med} entries

""", unsafe_allow_html=True, ) st.set_page_config( page_title="BioRxiv Search", page_icon=":scroll:", ) define_style() df, embeddings_unique = load_data_embeddings() logo(df["date"].max(), df[df['server'] == 'biorxiv'].shape[0], df[df['server'] == 'medrxiv'].shape[0]) # model = model_to_device() corpus_index = None corpus_precision = "ubinary" use_hf = False query = st.text_input("Enter your search query:") num_to_show = st.number_input( "Top k results to display:", min_value=1, max_value=50, value=10, ) st.markdown(""" """, unsafe_allow_html=True) search_button = st.button("Search") # 搜索逻辑触发 if query or search_button: with st.spinner("Searching..."): # Encode the query search_start_time = time.time() # query_embedding = model.encode([query], normalize_embeddings=True, precision=corpus_precision) embedding_time = time.time() raw_embadding = query_hf_api(query) query_embedding = process_embeddings(raw_embadding) embedding_time_total = time.time() - embedding_time # Perform the search results, search_time, corpus_index = semantic_search_faiss( query_embedding, corpus_index=corpus_index, corpus_embeddings=embeddings_unique if corpus_index is None else None, corpus_precision=corpus_precision, top_k=num_to_show, # type: ignore calibration_embeddings=None, rescore=False, rescore_multiplier=4, exact=True, output_index=True, ) search_end_time = time.time() search_duration = search_end_time - search_start_time st.markdown( f"
Search Completed in {search_duration:.2f} seconds (embeddings time: {embedding_time_total:.2f})
", unsafe_allow_html=True, ) # Prepare the results for plotting plot_data = {"Date": [], "Title": [], "Score": [], "DOI": [], "category": [], "server": []} search_df = pd.DataFrame(results[0]) # Find the minimum and maximum original scores min_score = search_df["score"].min() max_score = search_df["score"].max() # Normalize scores. The best score (min_score) becomes 100%, and the worst score (max_score) gets a value above 0%. search_df["score"] = abs(search_df["score"] - max_score) + min_score abstracts = [] # Iterate over each row in the search_df DataFrame for index, entry in search_df.iterrows(): row = df.iloc[int(entry["corpus_id"])] # Construct the DOI link try: doi_link = f"{doi.get_real_url_from_doi(row['doi'])}" except: doi_link = f'https://www.doi.org/' + row['doi'] # Append information to plot_data for visualization plot_data["Date"].append(row["date"]) plot_data["Title"].append(row["title"]) plot_data["Score"].append(search_df["score"][index]) # type: ignore plot_data["DOI"].append(row["doi"]) plot_data["category"].append(row["category"]) plot_data["server"].append(row["server"]) with st.expander(f"{index + 1}\. {row['title']}"): # type: ignore col1, col2 = st.columns(2) col1.markdown(f"**Score:** {entry['score']:.1f}") col2.markdown(f"**Server:** [{row['server']}]") st.markdown(f"**Authors:** {row['authors']}") col1, col2 = st.columns(2) col2.markdown(f"**Category:** {row['category']}") col1.markdown(f"**Date:** {row['date']}") # st.markdown(f"**Summary:**\n{summary_text}", unsafe_allow_html=False) abstracts.append(row['abstract']) st.markdown( f"**Abstract:**\n{row['abstract']}", unsafe_allow_html=False ) st.markdown( f"**[Full Text Read]({doi_link})** 🔗", unsafe_allow_html=True ) summary_button_one_paper = st.button("AI summary of this Paper", key="b_"+str(index+1)) if summary_button_one_paper: with st.spinner("AI summary of this Paper..."): ai_gen_start = time.time() st.markdown('**AI summary of this Paper:**') summary_of_this_Paper = summarize_abstract_kimi(title=row['title'], link=doi_link) st.markdown(summary_of_this_Paper) new_link = f"https://kimi.moonshot.cn/_prefill_chat?prefill_prompt=The paper we are going to discuss is {row['title']}, the link is {str(doi_link)} or https://www.{str(row['server'])}.org/content/{str(row['doi'])}v1 " \ f" Please use this as a basis to continue summarize this article and answer my follow-up questions &send_immediately=true&force_search=false" total_ai_time = time.time() - ai_gen_start st.markdown(f'**Time to generate summary:** {total_ai_time:.2f} s') # Make sure the HTML link is formatted correctly st.markdown(f'**Full Text Dialogue** 🔗', unsafe_allow_html=True) if plot_data: with st.spinner("Under statistics..."): plot_df = pd.DataFrame(plot_data) # Convert 'Date' to datetime if it's not already in that format plot_df["Date"] = pd.to_datetime(plot_df["Date"]) # Sort the DataFrame based on the Date to make sure it's ordered plot_df = plot_df.sort_values(by="Date") # Create a Plotly figure fig = px.scatter( plot_df, x="Date", y="Score", hover_data=["Title", "DOI"], color='server', title="Publication Times and Scores", ) fig.update_traces(marker=dict(size=10)) # Customize hover text to display the title and link it to the DOI fig.update_traces( hovertemplate="%{hovertext}", hovertext=plot_df.apply(lambda row: f"{row['Title']}", axis=1), ) # Show the figure in the Streamlit app st.plotly_chart(fig, use_container_width=True) # Generate category counts for the pie chart category_counts = plot_df["category"].value_counts().reset_index() category_counts.columns = ["category", "count"] # Create a pie chart with Plotly Express fig = px.pie( category_counts, values="count", names="category", title="Category Distribution", ) # Show the pie chart in the Streamlit app st.plotly_chart(fig, use_container_width=True) if abstracts: with st.spinner("LLM is summarizing..."): prompt = st.text_area("Enter your summary prompt", value=LLM_prompt) summary_button = st.button("AI summary", key="b2") if summary_button: ai_gen_start = time.time() st.markdown('**AI Summary of 10 abstracts:**') st.markdown(summarize_abstract(abstracts[:9], instructions=prompt)) total_ai_time = time.time() - ai_gen_start st.markdown(f'**Time to generate summary:** {total_ai_time:.2f} s')