Spaces:
Sleeping
Sleeping
import gradio as gr | |
import numpy as np | |
import os | |
import requests | |
from transformers import pipeline, AutoModelForTokenClassification, AutoTokenizer | |
from sentence_transformers import SentenceTransformer | |
from typing import List | |
NER_MODEL_PATH = 'dell-research-harvard/historical_newspaper_ner' | |
EMBED_MODEL_PATH = 'dell-research-harvard/same-story' | |
AZURE_VM_ALABAMA = os.environ.get('AZURE_VM_ALABAMA') | |
def find_sep_token(tokenizer): | |
""" | |
Returns sep token for given tokenizer | |
""" | |
if 'eos_token' in tokenizer.special_tokens_map: | |
sep = " " + tokenizer.special_tokens_map['eos_token'] + " " + tokenizer.special_tokens_map['sep_token'] + " " | |
else: | |
sep = " " + tokenizer.special_tokens_map['sep_token'] + " " | |
return sep | |
def find_mask_token(tokenizer): | |
""" | |
Returns mask token for given tokenizer | |
""" | |
mask_tok = tokenizer.special_tokens_map['mask_token'] | |
return mask_tok | |
if gr.NO_RELOAD: | |
ner_model=AutoModelForTokenClassification.from_pretrained(NER_MODEL_PATH) | |
ner_tokenizer=AutoTokenizer.from_pretrained(NER_MODEL_PATH, return_tensors = "pt", | |
max_length=256, truncation = True) | |
token_classifier = pipeline(task = "ner", | |
model = ner_model, tokenizer = ner_tokenizer, | |
ignore_labels = [], aggregation_strategy='max') | |
embedding_tokenizer = AutoTokenizer.from_pretrained(EMBED_MODEL_PATH) | |
embedding_model = SentenceTransformer(EMBED_MODEL_PATH) | |
embed_mask_tok = find_mask_token(embedding_tokenizer) | |
embed_sep_tok = find_sep_token(embedding_tokenizer) | |
# with open(REF_INDEX_PATH, 'r') as f: | |
# news_paths = [l.strip() for l in f.readlines()] | |
def handle_punctuation_for_generic_mask(word): | |
"""If punctuation comes before the word, return it before the mask, ow return it after the mask""" | |
if word[0] in [".",",","!","?"]: | |
return word[0] + " [MASK]" | |
elif word[-1] in [".",",","!","?"]: | |
return "[MASK]" + word[-1] | |
else: | |
return "[MASK]" | |
def handle_punctuation_for_entity_mask(word,entity_group): | |
"""If punctuation comes before the word, return it before the mask, ow return it after the mask - this is for specific entity masks""" | |
if word[0] in [".",",","!","?"]: | |
return word[0]+" "+entity_group | |
elif word[-1] in [".",",","!","?"]: | |
return entity_group+word[-1] | |
else: | |
return entity_group | |
def replace_words_with_entity_tokens(ner_output_dict: List[dict], | |
desired_labels: List[str] = ['PER', 'ORG', 'LOC', 'MISC'], | |
all_masks_same: bool = True) -> str: | |
if not all_masks_same: | |
new_word_list=[subdict["word"] if subdict["entity_group"] not in desired_labels else handle_punctuation_for_entity_mask(subdict["word"],subdict["entity_group"]) for subdict in ner_output_dict] | |
else: | |
new_word_list=[subdict["word"] if subdict["entity_group"] not in desired_labels else handle_punctuation_for_generic_mask(subdict["word"]) for subdict in ner_output_dict] | |
return " ".join(new_word_list) | |
def mask(ner_output_list: List[List[dict]], desired_labels: List[str] = ['PER', 'ORG', 'LOC', 'MISC'], | |
all_masks_same: bool = True) -> List[str]: | |
return replace_words_with_entity_tokens(ner_output_list, desired_labels, all_masks_same) | |
def ner(text: List[str]) -> List[str]: | |
results = token_classifier(text) | |
return results[0] | |
def ner_and_mask(text: List[str], labels_to_mask: List[str] = ['PER', 'ORG', 'LOC', 'MISC'], all_masks_same: bool = True) -> List[str]: | |
ner_output_list = ner(text) | |
return mask(ner_output_list, labels_to_mask, all_masks_same) | |
def embed(text: str) -> List[str]: | |
data = [] | |
# Correct [MASK] token for tokenizer | |
text = text.replace('[MASK]', embed_mask_tok) | |
text = text.replace('[SEP]', embed_sep_tok) | |
data.append(text) | |
embedding = embedding_model.encode(data, show_progress_bar = False, batch_size = 1) | |
embedding = embedding / np.linalg.norm(embedding, axis = 1, keepdims = True) | |
return embedding | |
def query(sentence: str) -> List[str]: | |
mask_results = ner_and_mask([sentence]) | |
embedding = embed(mask_results) | |
assert embedding.shape == (1, 768) | |
embedding = embedding[0].astype(np.float64) | |
req = {"vector": list(embedding), 'nn': 5} | |
# Send embedding to Azure VM | |
response = requests.post(f"http://{AZURE_VM_ALABAMA}/retrieve", json = req) | |
doc = response.json() | |
article = doc['bboxes'][doc['article_id']] | |
return article['raw_text'] | |
if __name__ == "__main__": | |
demo = gr.Interface( | |
fn=query, | |
inputs=["text"], | |
outputs=["text"], | |
) | |
demo.launch() |