import knowledge_extraction model

app.py

@@ -1,20 +1,19 @@
 import streamlit as st
-from transformers import BertTokenizer, BertForSequenceClassification
-import torch
+from transformers import pipeline
 
-# Load the pre-trained BERT model and tokenizer
-model_name = "bert-base-uncased"  # Replace with your specific model if needed
-tokenizer = BertTokenizer.from_pretrained(model_name)
-model = BertForSequenceClassification.from_pretrained(model_name)
+# Load the token classification pipeline
+model_name = "jjzha/jobbert_knowledge_extraction"
+pipe = pipeline("token-classification", model=model_name)
 
 # Streamlit UI
-st.title("BERT Text Classification")
-text_input = st.text_input("Enter text for classification:")
+st.title("Token Classification with Hugging Face")
+text_input = st.text_area("Enter text for token classification:")
 
 if text_input:
-    # Tokenize input
-    inputs = tokenizer(text_input, return_tensors="pt", truncation=True, padding=True, max_length=512)
-    with torch.no_grad():
-        logits = model(**inputs).logits
-    predicted_class = torch.argmax(logits, dim=1).item()
-    st.write(f"Predicted Class: {predicted_class}")
+    # Perform token classification
+    results = pipe(text_input)
+
+    # Display the results
+    st.write("Token Classification Results:")
+    for result in results:
+        st.write(f"Entity: {result['entity']}, Word: {result['word']}, Score: {result['score']:.2f}")