File size: 1,012 Bytes
9a75d7c
6596636
6dbe383
ba093b8
 
8fe4b6c
bcc7480
c447bbf
a38d800
 
6dbe383
a38d800
6dbe383
89cd9d2
28f0c2b
ba093b8
50a2b09
6dbe383
 
 
1886f29
bafdc64
c5042ea
50979de
 
 
c5042ea
 
50979de
 
c5042ea
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
from transformers import pipeline, AutoTokenizer
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:")
        tuples = ""
        # Process each entity and highlight the labeled words
        for entity in response:
            label = entity['entity']
            word =  entity["word"]
            tuples += f"{word}\t{label}\n"

        # Display the highlighted text using st.markdown
        st.text(tuples)