Update app.py

app.py

@@ -1,6 +1,20 @@
-from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
+import streamlit as st
+from transformers import BertTokenizer, BertForSequenceClassification
 import torch
 
-# Load the tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("jjzha/jobbert_knowledge_extraction")
-model = AutoModelForTokenClassification.from_pretrained("jjzha/jobbert_knowledge_extraction")
+# 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)
+
+# Streamlit UI
+st.title("BERT Text Classification")
+text_input = st.text_input("Enter text for 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}")