Update app.py

app.py

@@ -68,7 +68,8 @@ class ByT5ForTextGeotagging(PreTrainedModel):
         if return_coordinates:
             class_idx = torch.argmax(logits, dim=1).item()
             coordinates = self.config.class_to_location.get(str(class_idx))
-            return logits, coordinates
+            confidence = torch.max(logits).item()
+            return logits, coordinates, confidence
         else:
             return logits
 
@@ -84,8 +85,8 @@ byt5_tokenizer, model = load_model_and_tokenizer()
 
 def geolocate_text_byt5(text):
     input_tensor = byt5_tokenizer(text, return_tensors="pt", truncation=True, max_length=140)['input_ids']
-    logits, (lat, lon) = model(input_tensor.unsqueeze(0), return_coordinates=True)
-    return lat, lon
+    logits, (lat, lon), confidence = model(input_tensor.unsqueeze(0), return_coordinates=True)
+    return lat, lon, confidence
 
 if 'text_input' not in st.session_state:
     st.session_state.text_input = ""
@@ -123,8 +124,8 @@ if st.button('Submit'):
 
 
 if st.session_state.text_input:
-    location = geolocate_text_byt5(st.session_state.text_input)
-    st.write('Predicted Location: ', location)
+    location, confidence = geolocate_text_byt5(st.session_state.text_input)
+    st.write('Predicted Location: ', location, 'confidence: ', 'High' if confidence > 0.2 else 'Low')
 
      # Render map with pydeck
     map_data = pd.DataFrame(