app.py

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(
        """
    <style>
        .stExpander > .stButton > button {
            width: 100%;
            border: none;
            background-color: #f0f2f6;
            color: #333;
            text-align: left;
            padding: 15px;
            font-size: 18px;
            border-radius: 10px;
            margin-top: 5px;
        }
        .stExpander > .stExpanderContent {
            padding-left: 10px;
            padding-top: 10px;
        }
        a {
            color: #FF4B4B;
            text-decoration: none;
        }
    </style>
    """,
        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"""
        <div style='display: flex; justify-content: center; align-items: center;'>
            <div style='margin-right: 20px;'>
                <img src='{biorxiv_logo}' alt='BioRxiv logo' style='max-height: 100px;'>
            </div>
            <div style='margin-left: 20px;'>
                <img src='{medarxiv_logo}' alt='medRxiv logo' style='max-height: 100px;'>
            </div>
        </div>
        <div style='text-align: center; margin-top: 10px;'>
            <h3 style='color: black;'>LLM-based search and summary tool for bioRxiv & medRxiv</h3>
        </div>
        <p>How to use:
        <br><strong>1</strong>: Enter your search query (Optional modification "Top k results to display") 
        <br><strong>2</strong>: Press Enter in the query box (or click the search button) to search.
        <br><strong>3</strong>: When the search results are displayed, you can click on the line of interest to view the overview information of this article.
        <br><strong>4</strong>: If you want to learn more about the paper, you can jump to the paper pdf by clicking on 'Full Text Read' link.
        <br><strong>5</strong>: Enter summary prompt in the prompt input box.
        <br><strong>6</strong>: Click "AI summary" to summarize the search results above.
        </p>
        <div style='text-align: left; margin-top: 10px;'>
            Last database update: {db_update_date}; Database size: bioRxiv: {db_size_bio} / medRxiv: {db_size_med} entries
        </div>
        <br>
        """,
        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("""
    <style>
    div.stButton > button {
        width: 710px; /* 设置按钮宽度 */
        background-color: #007BFF; /* 蓝色背景 */
        color: white; /* 按钮文字颜色 */
        border: none; /* 去除边框 */
        border-radius: 5px; /* 圆角按钮 */
        padding: 10px; /* 内边距调整 */
        cursor: pointer; /* 鼠标悬浮样式 */
    }
    div.stButton > button:hover {
        background-color: #0056b3; /* 悬浮时更深的蓝色 */
    }
    </style>
    """, 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"<h6 style='text-align: center; color: #7882af;'>Search Completed in {search_duration:.2f} seconds (embeddings time: {embedding_time_total:.2f})</h3>",
            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'<a href="{new_link}" target="_blank">**Full Text Dialogue** 🔗</a>',
                                    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="<b>%{hovertext}</b>",
                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')