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