Update app.py

app.py

@@ -45,14 +45,15 @@ st.set_page_config(
     }
 )
 
-aspects = [
-    'task', 'method', 'dataset'
-]
+aspect_labels = {
+    'task': 'Task 🎯 ',
+    'method': 'Method 🔨 ',
+    'dataset': 'Dataset 🏷️',
+}
+aspects = list(aspect_labels.keys())
 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):
@@ -63,7 +64,7 @@ def st_load_model(name_or_path):
 
 @st.cache(show_spinner=False)
 def st_load_dataset(name_or_path):
-    with st.spinner('Loading the dataset (this might take a while)...'):
+    with st.spinner('Loading the dataset and search index (this might take a while)...'):
         dataset = load_dataset(name_or_path)
 
         if isinstance(dataset, DatasetDict):
@@ -84,6 +85,7 @@ aspect_to_model = dict(
 dataset = st_load_dataset(dataset_config)
 
 
+@st.cache(show_spinner=False)
 def get_paper(doc_id):
     res = requests.get(f'https://api.semanticscholar.org/v1/paper/{doc_id}')
 
@@ -93,32 +95,35 @@ def get_paper(doc_id):
         raise ValueError(f'Cannot load paper from S2 API: {doc_id}')
 
 
+def get_embedding(input_text, user_aspect):
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path)
+
+    # preprocess the input
+    inputs = tokenizer(input_text, padding=True, truncation=True, return_tensors="pt", max_length=512)
+
+    # inference
+    outputs = aspect_to_model[user_aspect](**inputs)
+
+    # 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)
+
+    return embeddings.detach().numpy()[0]
+
+
+@st.cache(show_spinner=False)
 def find_related_papers(paper_id, user_aspect):
     with st.spinner('Searching for related papers...'):
+        paper_id = paper_id.strip()  # remove white spaces
 
         paper = get_paper(paper_id)
 
         if paper is None or 'title' not in paper or paper['title'] is None or 'abstract' not in paper or paper['abstract'] is None:
-            raise ValueError(f'Could not retrieve title and abstract for input paper: {paper_id}')
+            raise ValueError(f'Could not retrieve title and abstract for input paper (the paper is probably behind a paywall): {paper_id}')
 
         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,
@@ -129,7 +134,7 @@ def find_related_papers(paper_id, user_aspect):
         ))
 
         # Retrieval
-        prompt = embeddings.detach().numpy()[0]
+        prompt = get_embedding(title_abs, user_aspect)
         scores, retrieved_examples = dataset.get_nearest_examples(f'{user_aspect}_embeddings', prompt, k=10)
 
         result.update(dict(
@@ -144,9 +149,9 @@ 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. 
+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 dataset close 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.
+st.markdown("""Enter a ArXiv ID or a DOI of a paper for that you want find similar papers. The title and abstract of the input paper must be available through the [Semantic Scholar API](https://www.semanticscholar.org/product/api).
 
 Try it yourself! 👇""",
     unsafe_allow_html=True)
@@ -165,6 +170,8 @@ with st.form("aspect-input", clear_on_submit=False):
         "ACL:N19-1423": "BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding",
         "10.18653/v1/S16-1001": "SemEval-2016 Task 4: Sentiment Analysis in Twitter",
         "10.1145/3065386": "ImageNet classification with deep convolutional neural networks",
+        "arXiv:2101.08700": "Multi-sense embeddings through a word sense disambiguation process",
+        "10.1145/3340531.3411878": "Incremental and parallel computation of structural graph summaries for evolving graphs",
     }
 
     example = st.selectbox(
@@ -175,7 +182,8 @@ with st.form("aspect-input", clear_on_submit=False):
 
     user_aspect = st.radio(
         label="In what aspect are you interested?",
-        options=aspects
+        options=aspects,
+        format_func=lambda option_key: aspect_labels[option_key],
     )
 
     cols = st.columns(3)