Update app.py

app.py

@@ -82,6 +82,27 @@ def load_model_and_tokenizer():
 
 byt5_tokenizer, model = load_model_and_tokenizer()
 
+def geolocate_text_byt5_multiclass(text):
+    input_tensor = byt5_tokenizer(text, return_tensors="pt", truncation=True, max_length=140)['input_ids']
+    logits, (lat, lon), confidence = model(input_tensor.unsqueeze(0), return_coordinates=True)
+    probas = torch.nn.functional.softmax(logits, dim=1).detach().cpu().numpy()
+
+    # Sort probabilities in descending order and get their indices
+    sorted_indices = np.argsort(-probas[0])
+
+    results = []
+    cumulative_prob = 0.0
+    for class_idx in sorted_indices:
+        prob = probas[0][class_idx]
+        cumulative_prob += prob
+        if cumulative_prob > 0.5:
+            coordinates = model.config.class_to_location.get(str(class_idx))
+            if coordinates:
+                results.append((class_idx, prob, coordinates))
+            break
+
+    return results
+   
 
 def geolocate_text_byt5(text):
     input_tensor = byt5_tokenizer(text, return_tensors="pt", truncation=True, max_length=140)['input_ids']
@@ -124,7 +145,8 @@ if st.button('Submit'):
 
 
 if st.session_state.text_input:
-    lat, lon, confidence = geolocate_text_byt5(st.session_state.text_input)
+    results = geolocate_text_byt5_multiclass(st.session_state.text_input)
+    _, confidence, (lat, lon) = results[0]
     st.write('Predicted Location: (', lat, lon, '). Confidence: ', 'High' if confidence > 0.2 else 'Low')
 
      # Render map with pydeck