app.py

import gradio as gr
import spaces
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch.nn.functional as F
import torch.nn as nn
import re
import requests
from urllib.parse import urlparse
import xml.etree.ElementTree as ET

# Model repository and device setup
model_path = 'ssocean/NAIP'
device = 'cuda' if torch.cuda.is_available() else 'cpu'

# Globals
model = None
tokenizer = None

def fetch_arxiv_paper(arxiv_input):
    """Fetch title & abstract from arXiv URL or ID."""
    try:
        if 'arxiv.org' in arxiv_input:
            parsed = urlparse(arxiv_input)
            arxiv_id = parsed.path.split('/')[-1].replace('.pdf', '')
        else:
            arxiv_id = arxiv_input.strip()
        api_url = f'http://export.arxiv.org/api/query?id_list={arxiv_id}'
        resp = requests.get(api_url)
        if resp.status_code != 200:
            return {"title":"", "abstract":"", "success":False, "message":"arXiv API error"}
        root = ET.fromstring(resp.text)
        ns = {'atom':'http://www.w3.org/2005/Atom'}
        entry = root.find('.//atom:entry', ns)
        if entry is None:
            return {"title":"", "abstract":"", "success":False, "message":"Paper not found"}
        title = entry.find('atom:title', ns).text.strip()
        abstract = entry.find('atom:summary', ns).text.strip()
        return {"title":title, "abstract":abstract, "success":True, "message":"Fetched successfully"}
    except Exception as e:
        return {"title":"", "abstract":"", "success":False, "message":f"Error: {e}"}

@spaces.GPU(duration=60, enable_queue=True)
def predict(title, abstract):
    """Predict a normalized impact score (0–1) from title & abstract."""
    global model, tokenizer

    if model is None:
        # Try loading full-precision on GPU first
        try:
            model = AutoModelForSequenceClassification.from_pretrained(
                model_path,
                num_labels=1,
                torch_dtype=torch.float32,
                device_map="auto"
            )
        except RuntimeError:
            # Fallback to CPU-only
            model = AutoModelForSequenceClassification.from_pretrained(
                model_path,
                num_labels=1,
                torch_dtype=torch.float32,
                device_map="cpu"
            )
        tokenizer = AutoTokenizer.from_pretrained(model_path)
        model.eval()

    prompt = (
        f"Given a certain paper,\n"
        f"Title: {title.strip()}\n"
        f"Abstract: {abstract.strip()}\n"
        f"Predict its normalized academic impact (0~1):"
    )
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024)
    inputs = {k: v.to(device) for k, v in inputs.items()}
    with torch.no_grad():
        logits = model(**inputs).logits
    prob = torch.sigmoid(logits).item()
    score = min(1.0, prob + 0.05)  # +0.05 adjustment
    return round(score, 4)

def get_grade_and_emoji(score):
    if score >= 0.900: return "AAA 🌟"
    if score >= 0.800: return "AA ⭐"
    if score >= 0.650: return "A ✨"
    if score >= 0.600: return "BBB 🔵"
    if score >= 0.550: return "BB 📘"
    if score >= 0.500: return "B 📖"
    if score >= 0.400: return "CCC 📝"
    if score >= 0.300: return "CC ✏️"
    return "C 📑"

def validate_input(title, abstract):
    """Ensure title ≥3 words, abstract ≥50 words, and ASCII-only."""
    non_ascii = re.compile(r'[^\x00-\x7F]')
    if len(title.split()) < 3:
        return False, "Title must be at least 3 words."
    if len(abstract.split()) < 50:
        return False, "Abstract must be at least 50 words."
    if non_ascii.search(title):
        return False, "Title contains non-ASCII characters."
    if non_ascii.search(abstract):
        return False, "Abstract contains non-ASCII characters."
    return True, "Inputs look good."

def update_button_status(title, abstract):
    valid, msg = validate_input(title, abstract)
    if not valid:
        return gr.update(value="Error: " + msg), gr.update(interactive=False)
    return gr.update(value=msg), gr.update(interactive=True)

def process_arxiv_input(arxiv_input):
    if not arxiv_input.strip():
        return "", "", "Please enter an arXiv URL or ID"
    res = fetch_arxiv_paper(arxiv_input)
    if res["success"]:
        return res["title"], res["abstract"], res["message"]
    return "", "", res["message"]

css = """
.gradio-container {
    font-family: 'Arial', sans-serif;
}
.main-title {
    text-align: center;
    color: #2563eb;
    font-size: 2.5rem !important;
    margin-bottom: 1rem !important;
    background: linear-gradient(45deg, #2563eb, #1d4ed8);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
}
.input-section {
    background: #ffffff;
    padding: 2rem;
    border-radius: 1rem;
    box-shadow: 0 4px 6px rgba(0,0,0,0.1);
}
.result-section {
    background: #f8fafc;
    padding: 2rem;
    border-radius: 1rem;
    margin-top: 2rem;
}
.methodology-section, .example-section {
    background: #fff7ed;
    padding: 2rem;
    border-radius: 1rem;
    margin-top: 2rem;
}
.grade-display {
    font-size: 3rem;
    text-align: center;
    margin: 1rem 0;
}
.arxiv-input {
    background: #f3f4f6;
    padding: 1rem;
    border-radius: 0.5rem;
    margin-bottom: 1.5rem;
}
.arxiv-link {
    color: #2563eb;
    text-decoration: underline;
    font-size: 0.9em;
}
.arxiv-note {
    color: #666666;
    font-size: 0.9em;
    margin-top: 0.5em;
    margin-bottom: 0.5em;
}
"""

example_papers = [
    {
        "title": "Attention Is All You Need",
        "abstract": "The dominant sequence transduction models are based on complex recurrent or convolutional neural networks that include an encoder and a decoder. The best performing models also connect the encoder and decoder through an attention mechanism. We propose a new simple network architecture, the Transformer, based solely on attention mechanisms, dispensing with recurrence and convolutions entirely. Experiments on two machine translation tasks show these models to be superior in quality while being more parallelizable and requiring significantly less time to train.",
        "score": 0.982,
        "note": "Revolutionary paper introducing the Transformer architecture."
    },
    {
        "title": "Language Models are Few-Shot Learners",
        "abstract": "Recent work has demonstrated substantial gains on many NLP tasks and benchmarks by pre-training on a large corpus of text followed by fine-tuning on a specific task. While typically task-agnostic in architecture, this method still requires task-specific fine-tuning datasets of thousands or tens of thousands of examples. By contrast, humans can generally perform a new language task from only a few examples or from simple instructions—something which current NLP systems still largely struggle to do. Here we show that scaling up language models greatly improves task-agnostic, few-shot performance, sometimes even reaching competitiveness with prior state-of-the-art fine-tuning approaches.",
        "score": 0.956,
        "note": "Groundbreaking GPT-3 paper on few-shot learning."
    },
    {
        "title": "An Empirical Study of Neural Network Training Protocols",
        "abstract": "This paper presents a comparative analysis of different training protocols for neural networks across various architectures. We examine the effects of learning rate schedules, batch size selection, and optimization algorithms on model convergence and final performance. Our experiments span multiple datasets and model sizes, providing practical insights for deep learning practitioners.",
        "score": 0.623,
        "note": "Solid empirical comparison of training protocols."
    }
]

with gr.Blocks(theme=gr.themes.Default(), css=css) as iface:
    gr.Markdown("<div class='main-title'>Papers Impact: AI-Powered Research Impact Predictor</div>")
    gr.HTML("""
<a href="https://visitorbadge.io/status?path=https%3A%2F%2FVIDraft-PaperImpact.hf.space">
  <img src="https://api.visitorbadge.io/api/visitors?path=https%3A%2F%2FVIDraft-PaperImpact.hf.space&countColor=%23263759" />
</a>
""")

    with gr.Row():
        with gr.Column(elem_classes="input-section"):
            with gr.Group(elem_classes="arxiv-input"):
                gr.Markdown("### Import from arXiv")
                arxiv_input = gr.Textbox(
                    lines=1,
                    placeholder="e.g. 2504.11651",
                    label="arXiv URL or ID",
                    value="2504.11651"
                )
                gr.Markdown("""
                <p class="arxiv-note">
                  Click to use the default example or visit <a href="https://arxiv.org" class="arxiv-link" target="_blank">arxiv.org</a>
                </p>
                """)
                fetch_button = gr.Button("🔍 Fetch Paper Details", variant="secondary")

            gr.Markdown("### Or Enter Paper Details Manually")
            title_input = gr.Textbox(
                lines=2,
                placeholder="Enter paper title (minimum 3 words)...",
                label="Paper Title"
            )
            abstract_input = gr.Textbox(
                lines=5,
                placeholder="Enter paper abstract (minimum 50 words)...",
                label="Paper Abstract"
            )
            validation_status = gr.Textbox(label="Validation Status", interactive=False)
            submit_button = gr.Button("🎯 Predict Impact", interactive=False, variant="primary")

        with gr.Column(elem_classes="result-section"):
            score_output = gr.Number(label="Impact Score")
            grade_output = gr.Textbox(label="Grade", elem_classes="grade-display")

    with gr.Row(elem_classes="methodology-section"):
        gr.Markdown("""
### Scientific Methodology
- **Training Data**: Papers from CS.CV, CS.CL (NLP), and CS.AI fields  
- **Optimization**: NDCG optimization with Sigmoid activation & MSE loss  
- **Validation**: Cross-validated on historical citation data  
- **Architecture**: Transformer-based text encoder  
- **Metrics**: Citation-pattern analysis & research influence
        """)

    gr.Markdown("""
### Rating Scale
| Grade | Score Range | Description         | Emoji |
|-------|-------------|---------------------|-------|
| AAA   | 0.900–1.000 | Exceptional Impact  | 🌟    |
| AA    | 0.800–0.899 | Very High Impact    | ⭐    |
| A     | 0.650–0.799 | High Impact         | ✨    |
| BBB   | 0.600–0.649 | Above Average       | 🔵    |
| BB    | 0.550–0.599 | Moderate Impact     | 📘    |
| B     | 0.500–0.549 | Average Impact      | 📖    |
| CCC   | 0.400–0.499 | Below Average       | 📝    |
| CC    | 0.300–0.399 | Low Impact          | ✏️    |
| C     | <0.300      | Limited Impact      | 📑    |
    """)

    with gr.Row(elem_classes="example-section"):
        gr.Markdown("### Example Papers")
        for paper in example_papers:
            gr.Markdown(f"""
#### {paper['title']}
**Score**: {paper['score']} | **Grade**: {get_grade_and_emoji(paper['score'])}  
{paper['abstract']}  
*{paper['note']}*  
---
            """)

    # Event handlers
    title_input.change(update_button_status, [title_input, abstract_input], [validation_status, submit_button])
    abstract_input.change(update_button_status, [title_input, abstract_input], [validation_status, submit_button])
    fetch_button.click(process_arxiv_input, [arxiv_input], [title_input, abstract_input, validation_status])

    def run_prediction(t, a):
        s = predict(t, a)
        return s, get_grade_and_emoji(s)

    submit_button.click(run_prediction, [title_input, abstract_input], [score_output, grade_output])

if __name__ == "__main__":
    iface.launch()