vis-similarity SPECTER

vis-SPECTER.py

## run mongodb of files through SPECTER's API https://github.com/allenai/paper-embedding-public-apis
from typing import Dict, List
import json
import requests

URL = "https://model-apis.semanticscholar.org/specter/v1/invoke"
MAX_BATCH_SIZE = 16

def chunks(lst, chunk_size=MAX_BATCH_SIZE):
    """Splits a longer list to respect batch size"""
    for i in range(0, len(lst), chunk_size):
        yield lst[i : i + chunk_size]

def embed(papers):
    embeddings_by_paper_id: Dict[str, List[float]] = {}

    for chunk in chunks(papers):
        # Allow Python requests to convert the data above to JSON
        response = requests.post(URL, json=chunk)

        if response.status_code != 200:
            raise RuntimeError("Sorry, something went wrong, please try later!")

        for paper in response.json()["preds"]:
            embeddings_by_paper_id[paper["paper_id"]] = paper["embedding"]

    return embeddings_by_paper_id

### get data

import pymongo
## be sure to set user and pwd
client = pymongo.MongoClient("mongodb+srv://<user>:<pwd>@cluster0.c1tlg.mongodb.net/vissimilarity?retryWrites=true&w=majority")
db = client['vissimilarity']
col = db['docs']

r = db.docs.find({}, {'Title':1, 'Abstract':1})
json = list(r)

# hack to rename fields and convert _id to string

l = []

for i in json:
    i['paper_id'] = str(i['_id'])
    del i['_id']
    i['title'] = i['Title']
    del i['Title']
    i['abstract'] = i['Abstract']
    del i['Abstract']
    l.append(i)

# an example
l[0]

## run through the api
all_embeddings = embed(l)

import pandas as pd

df = pd.DataFrame(all_embeddings)