Update app.py

app.py

@@ -20,8 +20,9 @@ if user_input:
         response = pipe(contexted_ipnut)
         st.write("Response:")
         st.markdown(response)
-        # Sorting response by the 'start' position to handle overlapping entities correctly
-        response = sorted(response, key=lambda x: x['start'])
+         # Tokenize input to get token-level alignment
+        tokens = tokenizer(user_input, return_offsets_mapping=True, return_tensors='pt')
+        offset_mapping = tokens['offset_mapping'][0].numpy().tolist()
 
         # Initialize the highlighted text
         highlighted_text = ""
@@ -29,7 +30,7 @@ if user_input:
 
         # Process each entity and highlight the labeled words
         for entity in response:
-            start, end = entity['start'], entity['end']
+            start, end = offset_mapping[entity['index']][0], offset_mapping[entity['index']][1]
             label = entity['entity']
 
             # Add text before the entity
@@ -45,4 +46,4 @@ if user_input:
         highlighted_text += user_input[last_position:]
 
         # Display the highlighted text using st.markdown
-        st.markdown(highlighted_text, unsafe_allow_html=True)
+        st.markdown(highlighted_text, unsafe_allow_html=True)