app.py

import streamlit as st
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import fitz
import os

# Load the model and tokenizer
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:
        inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
        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

st.title("Paper Citation Classifier")

# Input text boxes for abstract, full text, and affiliations
abstract_input = st.text_area("Enter Abstract:")
full_text_input = st.text_area("Enter Full Text:")
affiliations_input = st.text_area("Enter Affiliations:")

# PDF upload option
uploaded_file = st.file_uploader("Upload a PDF file", type=["pdf"])
if uploaded_file is not None:
    file_text = extract_text_from_pdf(uploaded_file)
    st.text("Extracted Text from PDF:")
    st.text(file_text)

    # Concatenate inputs with [SEP]
    combined_text = f"{abstract_input} [SEP] {full_text_input} [SEP] {affiliations_input} [SEP] {file_text}"

    if st.button("Predict"):
        predicted_class = predict_class(combined_text)
        if predicted_class is not None:
            class_labels = ["Level 1", "Level 2", "Level 3", "Level 4"]
            st.text(f"Predicted Class: {class_labels[predicted_class]}")