Spaces:
Runtime error
Runtime error
| """ | |
| 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: https://arxiv.org/abs/2203.14541 | |
| - 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, | |
| } | |
| ) | |
| 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' | |
| 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 | |
| def st_load_dataset(name_or_path): | |
| with st.spinner('Loading the dataset and search index (this might take a while)...'): | |
| dataset = load_dataset(name_or_path) | |
| if isinstance(dataset, DatasetDict): | |
| dataset = dataset['train'] | |
| # load existing FAISS index for each aspect | |
| for a in aspects: | |
| dataset.load_faiss_index(f'{a}_embeddings', f'{a}_embeddings.faiss') | |
| 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 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 (the paper is probably behind a paywall): {paper_id}') | |
| title_abs = paper['title'] + ': ' + paper['abstract'] | |
| result = dict( | |
| paper=paper, | |
| aspect=user_aspect, | |
| ) | |
| result.update(dict( | |
| #embeddings=embeddings.tolist(), | |
| )) | |
| # Retrieval | |
| prompt = get_embedding(title_abs, user_aspect) | |
| scores, retrieved_examples = dataset.get_nearest_examples(f'{user_aspect}_embeddings', prompt, k=10) | |
| result.update(dict( | |
| related_papers=retrieved_examples, | |
| )) | |
| return result | |
| # Page | |
| st.title('Aspect-based Paper Similarity') | |
| st.markdown("""This demo showcases [Specialized Document Embeddings for Aspect-based Research Paper Similarity](https://arxiv.org/abs/2203.14541).""") | |
| # 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 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](https://arxiv.org/abs/2203.14541).""") | |
| 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) | |
| # 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_labels = { | |
| "arXiv:1902.06818": "Data augmentation for low resource sentiment analysis using generative adversarial networks", | |
| "arXiv:2202.06671": "Neighborhood Contrastive Learning for Scientific Document Representations with Citation Embeddings", | |
| "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( | |
| label='Or select an example:', | |
| options=list(example_labels.keys()), | |
| format_func=lambda option_key: f'{example_labels[option_key]} ({option_key})', | |
| ) | |
| user_aspect = st.radio( | |
| label="In what aspect are you interested?", | |
| options=aspects, | |
| format_func=lambda option_key: aspect_labels[option_key], | |
| ) | |
| cols = st.columns(3) | |
| submitted = cols[1].form_submit_button("Find related papers") | |
| # Listener | |
| if submitted: | |
| if paper_id or example: | |
| 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.') | |