app.py

from transformers import pipeline
import streamlit as st

@st.cache_resource
def context_text(text): return f"### Context\n{text}\n\n### Answer"

@st.cache_resource
def load_pipe():
    model_name = "MSey/pbt_CaBERT_7_c10731"
    return pipeline("token-classification", model=model_name), AutoTokenizer.from_pretrained(model_name)

pipe, tokenizer = load_pipe()

st.header("Test Environment for pbt_CaBERT_7_c10731")
user_input = st.text_input("Enter your Prompt here:", "")
contexted_ipnut = context_text(user_input)
context_len = len(contexted_ipnut)

if user_input:
    with st.spinner('Generating response...'):
        response = pipe(contexted_ipnut)
        st.write("Response:")
        st.markdown(response)
         # 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 = ""
        last_position = 0

        # Process each entity and highlight the labeled words
        for entity in response:
            start, end = offset_mapping[entity['index']][0], offset_mapping[entity['index']][1]
            label = entity['entity']

            # Add text before the entity
            highlighted_text += user_input[last_position:start]
            
            # Add the highlighted entity
            highlighted_text += f'<mark style="background-color: #FFFF00;">{user_input[start:end]}</mark><sup>{label}</sup>'

            # Update the last position
            last_position = end

        # Add remaining text after the last entity
        highlighted_text += user_input[last_position:]

        # Display the highlighted text using st.markdown
        st.markdown(highlighted_text, unsafe_allow_html=True)