Spaces:
Sleeping
Sleeping
File size: 4,739 Bytes
8b91654 f4d9dbf 8b91654 8fc3433 fa69787 8b91654 8fc3433 31fe8a6 8fc3433 2fb5dff 8fc3433 2fb5dff 8fc3433 2fb5dff dc6e8e4 edef27b dc6e8e4 2b84175 71418e0 edef27b 2b84175 2fb5dff 0bb86ad 2fb5dff 8fc3433 84135ec 8fc3433 2fb5dff 0bb86ad 2fb5dff a05144a 2fb5dff 7853cf0 |
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 117 118 119 120 |
import streamlit as st
from streamlit_pills import pills
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import fitz
import os
model = AutoModelForSequenceClassification.from_pretrained("REEM-ALRASHIDI/LongFormer-Paper-Citaion-Classifier")
tokenizer = AutoTokenizer.from_pretrained("allenai/longformer-base-4096")
def extract_text_from_pdf(file_path):
text = ''
with fitz.open(file_path) as pdf_document:
for page_number in range(pdf_document.page_count):
page = pdf_document.load_page(page_number)
text += page.get_text()
return text
def predict_class(text):
try:
max_length = 4096
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
uploaded_files_dir = "uploaded_files"
os.makedirs(uploaded_files_dir, exist_ok=True)
class_colors = {
0: "#d62728", # Level 1
1: "#ff7f0e", # Level 2
2: "#2ca02c", # Level 3
3: "#1f77b4" # Level 4
}
st.set_page_config(page_title="Paper Citation Classifier", page_icon="logo.png")
st.image("logo.png", width=70)
st.markdown('<div style="position: absolute; top: 5px; left: 5px;"></div>', unsafe_allow_html=True)
# col1, col2 = st.columns([1, 3])
st.title("Paper Citation Classifier")
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:")
options=["Nursing", "Physics", "Maths", "Chemical", "Nuclear", "Engineering" ,"Other"]
categories = pills("Select WoS category", options)
# categories = st.multiselect("Select WoS categories:", options)
combined_text = f"{title_input} [SEP] {abstract_input} [SEP] {full_text_input} [SEP] {affiliations_input} [SEP] {' [SEP] '.join(categories)}"
if st.button("Predict"):
with st.spinner("Predicting..."):
predicted_class = predict_class(combined_text)
if predicted_class is not None:
class_labels = ["Level 1 (Highly Cited Paper)", "Level 2 (Average Cited Paper)", "Level 3 (More Cited Paper)", "Level 4 (Low Cited Paper)"]
st.text("Predicted Class:")
for i, label in enumerate(class_labels):
if i == predicted_class:
st.markdown(
f'<div style="background-color: {class_colors[predicted_class]}; padding: 10px; border-radius: 5px; color: white; font-weight: bold;">{label}</div>',
unsafe_allow_html=True
)
else:
st.text(label)
elif option == "PDF":
uploaded_file = st.file_uploader("Upload a PDF file", type=["pdf"])
if uploaded_file is not None:
with st.spinner("Processing PDF..."):
file_path = os.path.join(uploaded_files_dir, uploaded_file.name)
with open(file_path, "wb") as f:
f.write(uploaded_file.getbuffer())
st.success("File uploaded successfully.")
st.text(f"File Path: {file_path}")
file_text = extract_text_from_pdf(file_path)
st.text("Extracted Text:")
st.text(file_text)
if st.button("Predict from PDF Text"):
with st.spinner("Predicting..."):
predicted_class = predict_class(file_text)
if predicted_class is not None:
class_labels = ["Level 1 (Highly Cited Paper)", "Level 2 (Average Cited Paper)", "Level 3 (More Cited Paper)", "Level 4 (Low Cited Paper)"]
st.text("Predicted Class:")
for i, label in enumerate(class_labels):
if i == predicted_class:
st.markdown(
f'<div style="background-color: {class_colors[predicted_class]}; padding: 10px; border-radius: 5px; color: white; font-weight: bold;">{label}</div>',
unsafe_allow_html=True
)
else:
st.text(label)
|