Updating UI

app.py

@@ -9,7 +9,16 @@ 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')
+
+AZURE_VMS = {}
+AVAILABLE_STATES = ['All States']
+for k, v in os.environ.items():
+    if 'AZURE_VM' in k:
+        AZURE_VMS[k.split('_')[-1]] = v
+        AVAILABLE_STATES.append(k.split('_')[-1].capitalize())
+
+AVAILABLE_YEARS = ['All Years']
+
 
 
 def find_sep_token(tokenizer):
@@ -113,26 +122,51 @@ def embed(text: str) -> List[str]:
 
     return embedding
 
-def query(sentence: str) -> List[str]:
+def query(sentence: str, state: str, years: List[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)
+    if state == 'All States':
+        pass
+    else:
+        vm_address = AZURE_VMS[state.upper()]
+        # Send embedding to Azure VM
+        response = requests.post(f"http://{vm_address}/retrieve", json = req)
+
     doc = response.json()
     article = doc['bboxes'][int(doc['article_id'])]
-    return article['raw_text']
+    
+    results = {
+        'newspaper_name': doc['lccn']['title'],
+        'location': doc['lccn']['dbpedia_ids'][0].replace('%2C_', ', '),
+        'date': doc['scan']['date'],
+        'article_text': article['raw_text'],
+        'pdf_link': doc['scan']['jp2_url'].replace('jp2', 'pdf')
+    }
+    
+    return results['newspaper_name'], results['location'], results['date'], results['article_text'], results['pdf_link']
     
 
 if __name__ == "__main__":
     demo = gr.Interface(
         fn=query,
-        inputs=["text"],
-        outputs=["text"],
+        inputs=[ 
+            gr.Textbox(lines=10, label="News Article"), 
+            gr.Dropdown(AVAILABLE_STATES, label="States to Search"),
+            gr.CheckboxGroup(AVAILABLE_YEARS, label="Years to Search")
+        ],
+        outputs=[
+            gr.Textbox(label="Newspaper Name"),
+            gr.Textbox(label="Location"), 
+            gr.Textbox(label="Date"),
+            gr.Textbox(lines = 10, label="Article Text OCR"),
+            gr.Textbox(label="PDF Link")
+        ]
     )
 
     demo.launch()