Update app.py

app.py

@@ -51,6 +51,12 @@ spatialDataset = PbfMapDataset(data_file_path = data_file_path,
 
 data_loader = DataLoader(spatialDataset, batch_size=1, num_workers=0, shuffle=False, pin_memory=False, drop_last=False) #issue needs to be fixed with num_workers not stopping after finished
 
+# Create a dictionary to map entity names to indices
+entity_index_dict = {entity['pivot_name']: i for i, entity in enumerate(spatialDataset)}
+
+# Ensure names are stored in lowercase for case-insensitive matching
+entity_index_dict = {name.lower(): index for name, index in entity_index_dict.items()}
+
 #Pre-aquire the SpaBERT embeddings for all geo-entities within our dataset
 def process_entity(batch, model, device):
     input_ids = batch['masked_input'].to(device)
@@ -90,9 +96,9 @@ for i, batch in enumerate(data_loader):
     spaBERT_embedding, input_ids = process_entity(batch, spaBERT_model, device)
     spaBERT_embeddings.append(spaBERT_embedding)
 
-st.write("SpaBERT Embedding shape:", spaBERT_embedding[0].shape)
-st.write("SpaBERT Embedding:", spaBERT_embedding[0])
-
+#st.write("SpaBERT Embedding shape:", spaBERT_embedding[0].shape)
+#st.write("SpaBERT Embedding:", spaBERT_embedding[0])
+embedding_cache = {}
 
 
 #Get BERT Embedding for review
@@ -108,14 +114,34 @@ def get_bert_embedding(review_text):
     bert_embedding = outputs.last_hidden_state[:, 0, :].detach()     #CLS Token
     return bert_embedding
 
-
-
-
-
-
-
-
-
+def get_spaBert_embedding(entity):
+    entity_index = entity_index_dict.get(entity.lower(), None)
+    return spaBERT_embeddings[entity_index]
+
+def processSpatialEntities(review, nlp):
+    doc = nlp(review)
+    entity_spans = [(ent.start, ent.end, ent.text, ent.label_) for ent in doc.ents]
+    token_embeddings = []
+
+    # Iterate over each entity span and process only geo entities
+    for start, end, text, label in entity_spans:
+        if label in ['FAC', 'ORG', 'LOC', 'GPE']:  # Filter to geo-entities
+            spaBert_emb = get_spaBert_embedding(text)
+            token_embeddings.append((text, spaBert_emb))
+    return token_embeddings
+
+# Function to read reviews from a text file
+def load_reviews_from_file(file_path):
+    reviews = {}
+    try:
+        with open(file_path, 'r', encoding='utf-8') as file:
+            for i, line in enumerate(file):
+                line = line.strip()
+                if line:  # Ensure the line is not empty
+                    reviews[f"Review {i + 1}"] = line
+    except FileNotFoundError:
+        st.error(f"File not found: {file_path}")
+    return reviews
 
 
 st.title("SpaGAN Demo")
@@ -138,10 +164,12 @@ for label, (color, description) in COLOR_MAP.items():
 #user_input = st.text_area("Input Text", height=200)
 
 # Define example reviews for testing
-example_reviews = {
-    "Review 1": "I visited the Empire State Building in New York last summer, and it was amazing!",
-    "Review 2": "Google, headquartered in Mountain View, is a leading tech company in the United States.",
-}
+#example_reviews = {
+#    "Review 1": "Meh. My brother lives near the Italian Market in South Philly. I went for a visit. Luckily for me, my brother and his girlfriend are foodies. I was able to taste many different cuisines in Philly. Coming from San Francisco, there are places I don't go due to the tourist trap aura and the non-authenticity of it all (Fisherman’s Wharf, Chinatown, etc.). But when I was in Philly, I had to have a cheesesteak... and I had to go to the two most famous places, which of course are right across the street from one another, in a big rivalry, and featured on the Food Network! How cheesy, but essential. We split two, both "wit whiz"? (cheese whiz) one from Geno's and one from Pat's. Pat's was much tastier than Geno's. The meat was seasoned, and the bun and cheese had much more flavor... better of the two... it seems.",
+#    "Review 2": "Google, headquartered in Mountain View, is a leading tech company in the United States.",
+#}
+review_file_path = "models/spabert/datasets/SampleReviews.txt"
+example_reviews = load_reviews_from_file(review_file_path)
 
 # Dropdown for selecting an example review
 user_input = st.selectbox("Select an example review", options=list(example_reviews.keys()))