Upload app.py

app.py

@@ -0,0 +1,254 @@
+"""
+
+Run via: streamlit run app.py
+
+"""
+
+import json
+import logging
+
+import requests
+import streamlit as st
+import torch
+from datasets import load_dataset
+from datasets.dataset_dict import DatasetDict
+from transformers import AutoTokenizer, AutoModel
+
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+    level=logging.INFO,
+)
+logger = logging.getLogger(__name__)
+
+model_hub_url = 'https://huggingface.co/malteos/aspect-scibert-task'
+
+about_page_markdown = f"""# 🔍 Find Papers With Similar Task
+
+See 
+- GitHub: https://github.com/malteos/aspect-document-embeddings
+- Paper: #TODO
+- Model hub: https://huggingface.co/malteos/aspect-scibert-task
+
+"""
+
+# Page setup
+st.set_page_config(
+    page_title="Papers with similar Task",
+    page_icon="🔍",
+    layout="centered",
+    initial_sidebar_state="auto",
+    menu_items={
+        'Get help': None,
+        'Report a bug': None,
+        'About': about_page_markdown,
+    }
+)
+
+aspects = [
+    'task', 'method', 'dataset'
+]
+tokenizer_name_or_path = f'malteos/aspect-scibert-{aspects[0]}'  # any aspect
+dataset_config = 'malteos/aspect-paper-metadata'
+
+tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path)
+
+
+@st.cache(show_spinner=False)
+def st_load_model(name_or_path):
+    with st.spinner(f'Loading the model `{name_or_path}` (this might take a while)...'):
+        model = AutoModel.from_pretrained(name_or_path)
+    return model
+
+
+@st.cache(show_spinner=False)
+def st_load_dataset(name_or_path):
+    with st.spinner('Loading the dataset (this might take a while)...'):
+        dataset = load_dataset(name_or_path)
+
+        if isinstance(dataset, DatasetDict):
+            dataset = dataset['train']
+
+        # load existing faiss
+        for a in aspects:
+            dataset.load_faiss_index(f'{a}_embeddings', f'{a}_embeddings.faiss')
+
+        # add faiss
+        #dataset.add_faiss_index(column=f'{aspect}_embeddings')
+        #loaded_dataset.add_faiss_index(column='method_embeddings')
+        #loaded_dataset.add_faiss_index(column='dataset_embeddings')
+
+    return dataset
+
+
+aspect_to_model = dict(
+    task=st_load_model('malteos/aspect-scibert-task'),
+    method=st_load_model('malteos/aspect-scibert-method'),
+    dataset=st_load_model('malteos/aspect-scibert-dataset'),
+)
+dataset = st_load_dataset(dataset_config)
+
+
+def get_paper(doc_id):
+    res = requests.get(f'https://api.semanticscholar.org/v1/paper/{doc_id}')
+
+    if res.status_code == 200:
+        return res.json()
+    else:
+        raise ValueError(f'Cannot load paper from S2 API: {doc_id}')
+
+
+def find_related_papers(paper_id, user_aspect):
+    # Add result to session
+
+    paper = get_paper(paper_id)
+
+    if paper is None or 'title' not in paper or 'abstract' not in paper:
+        raise ValueError('Could not retrieve data for input paper')
+
+    title_abs = paper['title'] + ': ' + paper['abstract']
+
+    # preprocess the input
+    inputs = tokenizer(title_abs, padding=True, truncation=True, return_tensors="pt", max_length=512)
+
+    # inference
+    outputs = aspect_to_model[user_aspect](**inputs)
+
+    # logger.info(f'attention_mask: {inputs["attention_mask"].shape}')
+    #
+    # logger.info(f'Outputs: {outputs["last_hidden_state"]}')
+    # logger.info(f'Outputs: {outputs["last_hidden_state"].shape}')
+
+    # Mean pool the token-level embeddings to get sentence-level embeddings
+    embeddings = torch.sum(
+        outputs["last_hidden_state"] * inputs['attention_mask'].unsqueeze(-1), dim=1
+    ) / torch.clamp(torch.sum(inputs['attention_mask'], dim=1, keepdims=True), min=1e-9)
+
+    result = dict(
+        paper=paper,
+        aspect=user_aspect,
+    )
+
+    result.update(dict(
+        #embeddings=embeddings.tolist(),
+    ))
+
+    # Retrieval
+    prompt = embeddings.detach().numpy()[0]
+    scores, retrieved_examples = dataset.get_nearest_examples(f'{user_aspect}_embeddings', prompt, k=10)
+
+    result.update(dict(
+        related_papers=retrieved_examples,
+    ))
+
+    # st.session_state.results.append(result)
+
+    return result
+
+
+# # Start session
+# if 'results' not in st.session_state:
+#     st.session_state.results = []
+
+# Page
+st.title('Aspect-based Paper Similarity')
+st.markdown("""This demo showcases [Specialized Document Embeddings for Aspect-based Research Paper Similarity](#TODO).""")
+
+# Introduction
+st.markdown(f"""The model was trained using a triplet loss on machine learning papers from the [paperswithcode.com](https://paperswithcode.com/) corpus with the objective of pulling embeddings of papers with the same task, method, or datasetclose together. For a more comprehensive overview of the model check out the [model card on 🤗 Model Hub]({model_hub_url}) or read [our paper](#TODO).
+""")
+st.markdown("""Enter a ArXiv ID or a DOI of a paper for that you want find similar papers.
+
+Try it yourself! 👇""",
+    unsafe_allow_html=True)
+
+# Demo
+with st.form("aspect-input", clear_on_submit=False):
+    paper_id = st.text_input(
+        label='Enter paper ID (format "arXiv:<arxiv_id>", "<doi>", or "ACL:<acl_id>"):',
+        # value="arXiv:2202.06671",
+        placeholder='Any DOI, ACL, or ArXiv ID'
+    )
+
+    example = st.selectbox(
+        label='Or select example',
+        options=[
+            "arXiv:2202.06671",
+            '10.1016/j.eswa.2019.06.026'
+        ]
+    )
+
+    # click_clear = st.button('clear text input', key=1)
+    # if click_clear:
+    #     paper_id = st.text_input(
+    #         label='Enter paper ID (arXiv:<arxiv_id>, or <doi>):', value="XXX", placeholder='123')
+
+    user_aspect = st.radio(
+        label="In what aspect are you interested?",
+        options=aspects
+    )
+
+    cols = st.columns(3)
+    submitted = cols[1].form_submit_button("Find related papers")
+
+# Listener
+if submitted:
+    if paper_id or example:
+        with st.spinner('Finding related papers...'):
+            try:
+                result = find_related_papers(paper_id if paper_id else example, user_aspect)
+
+                input_paper = result['paper']
+                related_papers = result['related_papers']
+
+                # with st.empty():
+
+                st.markdown(
+                    f'''Your input paper: \n\n<a href="{input_paper['url']}"><b>{input_paper['title']}</b></a> ({input_paper['year']})<hr />''',
+                    unsafe_allow_html=True)
+
+                related_html = '<ul>'
+
+                for i in range(len(related_papers['paper_id'])):
+                    related_html += f'''<li><a href="{related_papers['url_abs'][i]}">{related_papers['title'][i]}</a></li>'''
+
+                related_html += '</ul>'
+
+                st.markdown(f'''Related papers with similar {result['aspect']}: {related_html}''', unsafe_allow_html=True)
+
+            except (TypeError, ValueError, KeyError) as e:
+                st.error(f'**Error**: {e}')
+
+    else:
+        st.error('**Error**: No paper ID provided. Please provide a ArXiv ID or DOI.')
+
+# # Results
+# if 'results' in st.session_state and st.session_state.results:
+#     first = True
+#     for result in st.session_state.results[::-1]:
+#         if not first:
+#             st.markdown("---")
+#         # st.markdown(f"ID:\n> {result['paperId']}")
+#         # col_1, col_2, col_3 = st.columns([1,2,2])
+#         # col_1.metric(label='', value=json.dumps(result))
+#         # col_2.metric(label='Label', value=f"fooo")
+#         # col_3.metric(label='Score', value=f"123")
+#         input_paper = result['paper']
+#         related_papers = result['related_papers']
+#
+#         # with st.empty():
+#
+#         st.markdown(f'''Your input paper: \n\n<a href="{input_paper['url']}"><b>{input_paper['title']}</b></a> ({input_paper['year']})<hr />''', unsafe_allow_html=True)
+#
+#         related_html = '<ul>'
+#
+#         for i in range(len(related_papers['paper_id'])):
+#             related_html += f'''<li><a href="{related_papers['url_abs'][i]}">{related_papers['title'][i]}</a></li>'''
+#
+#         related_html += '</ul>'
+#
+#         st.markdown(f'''Related papers with similar {result['aspect']}: {related_html}''', unsafe_allow_html=True)
+#
+#         # st.markdown(f'''Related papers: {related_html}''', unsafe_allow_html=True)
+#
+#         first = False