Update app.py

app.py

@@ -9,29 +9,32 @@ import requests
 from urllib.parse import urlparse
 import xml.etree.ElementTree as ET
 
-model_path = r'ssocean/NAIP'
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
+# Model repository path and device selection
+model_path = "ssocean/NAIP"
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-global model, tokenizer
+# Global model/tokenizer variables
 model = None
 tokenizer = None
 
 def fetch_arxiv_paper(arxiv_input):
-    """Fetch paper details from arXiv URL or ID using requests."""
+    """
+    Fetch paper details (title, abstract) from an arXiv URL or ID using requests.
+    """
     try:
-        # Extract arXiv ID from URL or use directly
-        if 'arxiv.org' in arxiv_input:
+        # If user passed a full arxiv.org link, parse out the ID
+        if "arxiv.org" in arxiv_input:
             parsed = urlparse(arxiv_input)
             path = parsed.path
-            arxiv_id = path.split('/')[-1].replace('.pdf', '')
+            arxiv_id = path.split("/")[-1].replace(".pdf", "")
         else:
+            # Otherwise just use the raw ID
             arxiv_id = arxiv_input.strip()
 
-        # Fetch metadata using arXiv API
-        api_url = f'http://export.arxiv.org/api/query?id_list={arxiv_id}'
-        response = requests.get(api_url)
-        
-        if response.status_code != 200:
+        # ArXiv API query
+        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": "",
@@ -39,14 +42,10 @@ def fetch_arxiv_paper(arxiv_input):
                 "message": "Error fetching paper from arXiv API"
             }
 
-        # Parse the response XML
-        root = ET.fromstring(response.text)
-        
-        # ArXiv API uses Atom namespace
-        ns = {'arxiv': 'http://www.w3.org/2005/Atom'}
-        
-        # Extract title and abstract
-        entry = root.find('.//arxiv:entry', ns)
+        # Parse XML response
+        root = ET.fromstring(resp.text)
+        ns = {"arxiv": "http://www.w3.org/2005/Atom"}
+        entry = root.find(".//arxiv:entry", ns)
         if entry is None:
             return {
                 "title": "",
@@ -54,10 +53,9 @@ def fetch_arxiv_paper(arxiv_input):
                 "success": False,
                 "message": "Paper not found"
             }
-            
-        title = entry.find('arxiv:title', ns).text.strip()
-        abstract = entry.find('arxiv:summary', ns).text.strip()
-        
+
+        title = entry.find("arxiv:title", ns).text.strip()
+        abstract = entry.find("arxiv:summary", ns).text.strip()
         return {
             "title": title,
             "abstract": abstract,
@@ -74,34 +72,35 @@ def fetch_arxiv_paper(arxiv_input):
 
 @spaces.GPU(duration=60, enable_queue=True)
 def predict(title, abstract):
-    """Predict a normalized academic impact score (0–1) from title & abstract."""
-    title = title.replace("\n", " ").strip().replace("''", "'")
-    abstract = abstract.replace("\n", " ").strip().replace("''", "'")
+    """
+    Predict a normalized academic impact score (0–1) given the paper title & abstract.
+    Loads the model once globally, then uses it for inference.
+    """
     global model, tokenizer
 
     if model is None:
-        # Load config and disable any quantization
+        # Load model config, disable quantization, and set number of labels if needed
         config = AutoConfig.from_pretrained(model_path)
         config.quantization_config = None
+        config.num_labels = 1  # For classification/logit output
 
-        # Load model in full float32, then move to device
+        # IMPORTANT: Do not pass num_labels directly into from_pretrained for LLaMA-based models
         model = AutoModelForSequenceClassification.from_pretrained(
             model_path,
             config=config,
-            num_labels=1,
-            torch_dtype=torch.float32,
-            device_map=None,
+            torch_dtype=torch.float32,  # Use full-precision float32
+            device_map=None,           # We'll move it manually
             low_cpu_mem_usage=False
         )
         model.to(device)
+        model.eval()
 
         tokenizer = AutoTokenizer.from_pretrained(model_path)
-        model.eval()
-        
+
     text = (
         f"Given a certain paper,\n"
-        f"Title: {title}\n"
-        f"Abstract: {abstract}\n"
+        f"Title: {title.strip()}\n"
+        f"Abstract: {abstract.strip()}\n"
         f"Predict its normalized academic impact (0~1):"
     )
 
@@ -109,15 +108,17 @@ def predict(title, abstract):
         inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=1024)
         inputs = {k: v.to(device) for k, v in inputs.items()}
         with torch.no_grad():
-            outputs = model(**inputs)
-        prob = torch.sigmoid(outputs.logits).item()
-        score = min(1.0, prob + 0.05)
+            output = model(**inputs)
+        logits = output.logits
+        prob = torch.sigmoid(logits).item()
+        score = min(1.0, prob + 0.05)  # +0.05 offset, capped at 1.0
         return round(score, 4)
     except Exception as e:
         print(f"Prediction error: {e}")
-        return 0.0  # default on error
+        return 0.0  # Return 0 in case of any error
 
 def get_grade_and_emoji(score):
+    """Convert a 0–1 score into a tier grade with emoji indicator."""
     if score >= 0.900: return "AAA 🌟"
     if score >= 0.800: return "AA ⭐"
     if score >= 0.650: return "A ✨"
@@ -128,60 +129,41 @@ def get_grade_and_emoji(score):
     if score >= 0.300: return "CC ✏️"
     return "C 📑"
 
-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 that introduced the Transformer architecture, fundamentally changing NLP and deep learning."
-    },
-    {
-        "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 that demonstrated the power of large language models."
-    },
-    {
-        "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 research paper with useful findings but more limited scope and impact."
-    }
-]
-
 def validate_input(title, abstract):
-    title = title.replace("\n", " ").strip().replace("''", "'")
-    abstract = abstract.replace("\n", " ").strip().replace("''", "'")
-    non_latin_pattern = re.compile(r'[^\u0000-\u007F]')
+    """
+    Ensure title >=3 words, abstract >=50 words, and only ASCII chars.
+    """
+    non_ascii = re.compile(r"[^\x00-\x7F]")
     if len(title.split()) < 3:
-        return False, "The title must be at least 3 words long."
+        return False, "Title must be at least 3 words."
     if len(abstract.split()) < 50:
-        return False, "The abstract must be at least 50 words long."
-    if non_latin_pattern.search(title):
-        return False, "The title contains invalid characters. Only English letters and symbols are allowed."
-    if non_latin_pattern.search(abstract):
-        return False, "The abstract contains invalid characters. Only English letters and symbols are allowed."
-    return True, "Inputs are valid!"
+        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, message = validate_input(title, abstract)
+    valid, msg = validate_input(title, abstract)
     if not valid:
-        return gr.update(value="Error: " + message), gr.update(interactive=False)
-    return gr.update(value=message), gr.update(interactive=True)
+        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):
+    """
+    Helper to fill in title/abstract fields from an arXiv link/ID.
+    """
     if not arxiv_input.strip():
         return "", "", "Please enter an arXiv URL or ID"
     result = fetch_arxiv_paper(arxiv_input)
     if result["success"]:
         return result["title"], result["abstract"], result["message"]
-    else:
-        return "", "", result["message"]
+    return "", "", result["message"]
 
+# Custom CSS for styling
 css = """
-.gradio-container {
-    font-family: 'Arial', sans-serif;
-}
+.gradio-container { font-family: Arial, sans-serif; }
 .main-title {
     text-align: center;
     color: #2563eb;
@@ -201,7 +183,7 @@ css = """
     background: white;
     padding: 2rem;
     border-radius: 1rem;
-    box-shadow: 0 4px 6px -1px rgb(0 0 0 / 0.1);
+    box-shadow: 0 4px 6px -1px rgba(0,0,0,0.1);
 }
 .result-section {
     background: #f8fafc;
@@ -246,6 +228,53 @@ css = """
 }
 """
 
+# Example papers
+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 that introduced 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 that demonstrated the power of large language models."
+    },
+    {
+        "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 research paper with useful findings but more limited scope and impact."
+    }
+]
+
+# Build Gradio interface
 with gr.Blocks(theme=gr.themes.Default(), css=css) as iface:
     gr.Markdown(
         """
@@ -259,6 +288,7 @@ with gr.Blocks(theme=gr.themes.Default(), css=css) as iface:
 
     with gr.Row():
         with gr.Column(elem_classes="input-section"):
+            # arXiv import group
             with gr.Group(elem_classes="arxiv-input"):
                 gr.Markdown("### 📑 Import from arXiv")
                 arxiv_input = gr.Textbox(
@@ -267,16 +297,17 @@ with gr.Blocks(theme=gr.themes.Default(), css=css) as iface:
                     label="arXiv Paper URL/ID",
                     value="2504.11651"
                 )
-                gr.Markdown("""
-                <p class="arxiv-note">
-                Click input field to use example paper or browse papers at 
-                <a href="https://arxiv.org" target="_blank" class="arxiv-link">arxiv.org</a>
-                </p>
-                """)
+                gr.Markdown(
+                    """
+                    <p class="arxiv-note">
+                      Click input field to use example paper or browse papers at 
+                      <a href="https://arxiv.org" target="_blank" class="arxiv-link">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)...",
@@ -289,7 +320,7 @@ with gr.Blocks(theme=gr.themes.Default(), css=css) as iface:
             )
             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"):
             with gr.Group():
                 score_output = gr.Number(label="🎯 Impact Score")
@@ -330,7 +361,7 @@ with gr.Blocks(theme=gr.themes.Default(), css=css) as iface:
         for paper in example_papers:
             gr.Markdown(
                 f"""
-                #### {paper['title']} 
+                #### {paper['title']}
                 **Score**: {paper.get('score', 'N/A')} | **Grade**: {get_grade_and_emoji(paper.get('score', 0))}
                 {paper['abstract']}
                 *{paper['note']}*
@@ -338,6 +369,7 @@ with gr.Blocks(theme=gr.themes.Default(), css=css) as iface:
                 """
             )
 
+    # Validate button status on input changes
     title_input.change(
         update_button_status,
         inputs=[title_input, abstract_input],
@@ -348,13 +380,15 @@ with gr.Blocks(theme=gr.themes.Default(), css=css) as iface:
         inputs=[title_input, abstract_input],
         outputs=[validation_status, submit_button]
     )
-    
+
+    # Fetch from arXiv
     fetch_button.click(
         process_arxiv_input,
         inputs=[arxiv_input],
         outputs=[title_input, abstract_input, validation_status]
     )
 
+    # Predict callback
     def process_prediction(title, abstract):
         score = predict(title, abstract)
         grade = get_grade_and_emoji(score)