Inital commit for ndjv app

README.md

@@ -3,10 +3,10 @@ title: Newsdejavu
 emoji: 🌖
 colorFrom: yellow
 colorTo: green
-sdk: streamlit
+sdk: gradio
 sdk_version: 1.33.0
 app_file: app.py
-pinned: false
+pinned: falsesudo sudo 
 license: apache-2.0
 ---
 

app.py

@@ -1,6 +1,138 @@
-# Just trying out a comment to test commits
+import gradio as gr
+import numpy as np
+import os
+import requests
+from transformers import pipeline, AutoModelForTokenClassification, AutoTokenizer
+from sentence_transformers import SentenceTransformer
 
-import streamlit as st
+from typing import List
 
-x = st.slider('Select a value')
-st.write(x, 'squared is', x * x)
+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()

requirements.txt

@@ -0,0 +1,3 @@
+transformers
+sentenceistransformers
+numpy