File size: 4,930 Bytes
8e257e5 524c798 8e257e5 524c798 8e257e5 524c798 8e257e5 |
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 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
import streamlit as st
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import fitz
import os
# Load Longformer model and tokenizer
longformer_model = AutoModelForSequenceClassification.from_pretrained("Reem333/Longformer")
longformer_tokenizer = AutoTokenizer.from_pretrained("allenai/longformer-base-4096")
# Load BERT model and tokenizer
bert_model = AutoModelForSequenceClassification.from_pretrained("Reem333/BERT")
bert_tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
# Function to predict the class of the text using a specified model and tokenizer
def predict_class(text, model, tokenizer):
try:
max_length = 4096 if "longformer" in str(model) else 512
truncated_text = text[:max_length]
inputs = tokenizer(truncated_text, return_tensors="pt", padding=True, truncation=True, max_length=max_length)
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
predicted_class = torch.argmax(logits, dim=1).item()
return predicted_class
except Exception as e:
st.error(f"Error during prediction: {e}")
return None
# Setup for uploaded files directory
uploaded_files_dir = "uploaded_files"
os.makedirs(uploaded_files_dir, exist_ok=True)
# Color mapping for class levels
class_colors = {
0: "#d62728", # Level 1
1: "#ff7f0e", # Level 2
2: "#2ca02c", # Level 3
3: "#1f77b4" # Level 4
}
# Streamlit page configuration
st.set_page_config(page_title="Paper Citation Classifier", page_icon="logo.png")
# Sidebar content
with st.sidebar:
st.image("logo.png", width=70)
st.markdown('<div style="position: absolute; left: 5px;"></div>', unsafe_allow_html=True)
st.markdown("# Paper Citation Classifier")
st.markdown("---")
st.markdown("## About")
st.markdown('''
This tool classifies paper citations into different levels based on their number of citations.
Powered by Fine-Tuned [Longformer model](https://huggingface.co/Reem333/LongFormer) and BERT model with custom data.
''')
st.markdown("### Class Levels:")
st.markdown("- Level 1: Low cited papers")
st.markdown("- Level 2: Average cited papers")
st.markdown("- Level 3: More cited papers")
st.markdown("- Level 4: Highly cited papers")
st.markdown("---")
st.markdown('Tabuk University')
st.title("Check Your Paper Now!")
# Main content
option = st.radio("Select input type:", ("Text", "PDF"))
if option == "Text":
title_input = st.text_area("Enter Title:")
abstract_input = st.text_area("Enter Abstract:")
full_text_input = st.text_area("Enter Full Text:")
affiliations_input = st.text_area("Enter Affiliations:")
keywords_input = st.text_area("Enter Keywords:")
options = ['cs', "AI"]
selected_category = st.selectbox("Select WoS categories:", options)
if selected_category == "Other":
custom_category = st.text_input("Enter custom category:")
selected_category = custom_category if custom_category else "Other"
combined_text = f"{title_input} [SEP] {keywords_input} [SEP] {abstract_input} [SEP] {selected_category} [SEP] {affiliations_input} [SEP] {full_text_input}"
if st.button("Predict"):
if not any([title_input, abstract_input, keywords_input, full_text_input, affiliations_input]):
st.warning("Please enter paper text.")
else:
with st.spinner("Predicting..."):
longformer_class = predict_class(combined_text, longformer_model, longformer_tokenizer)
bert_class = predict_class(combined_text, bert_model, bert_tokenizer)
if longformer_class is not None and bert_class is not None:
class_labels = ["Level 1", "Level 2", "Level 3", "Level 4"]
st.text("Longformer Predicted Class:")
for i, label in enumerate(class_labels):
if i == longformer_class:
st.markdown(
f'<div style="background-color: {class_colors[longformer_class]}; padding: 10px; border-radius: 5px; color: white; font-weight: bold;">{label}</div>',
unsafe_allow_html=True
)
else:
st.text(label)
st.text("BERT Predicted Class:")
for i, label in enumerate(class_labels):
if i == bert_class:
st.markdown(
f'<div style="background-color: {class_colors[bert_class]}; padding: 10px; border-radius: 5px; color: white; font-weight: bold;">{label}</div>',
unsafe_allow_html=True
)
else:
st.text(label)
|