Upload app.py

app.py

@@ -14,14 +14,11 @@ from marker.logger import configure_logging
 from surya.settings import settings as surya_settings
 import traceback
 import re
-import time
 
-# =========================
-# Marker / PDF 预处理配置
-# =========================
+# marker
 configure_logging()
 MAX_PAGES = 30
-MIN_LENGTH = 200
+MIN_LENGTH=200
 settings.EXTRACT_IMAGES = False
 settings.DEBUG = False
 settings.PDFTEXT_CPU_WORKERS = 1
@@ -32,9 +29,7 @@ surya_settings.IN_STREAMLIT = True
 model_refs = load_all_models()
 metadata = {}
 
-# =========================
-# LLM 加载
-# =========================
+# prepare LLM
 model_name = "maxidl/Llama-OpenReviewer-8B"
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
@@ -42,48 +37,32 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="auto"
 )
 tokenizer = AutoTokenizer.from_pretrained(model_name)
+# Set pad_token to eos_token if not set (common for Llama models)
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
     tokenizer.pad_token_id = tokenizer.eos_token_id
+streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, decode_kwargs=dict(skip_special_tokens=True))
 
-streamer = TextIteratorStreamer(
-    tokenizer,
-    skip_prompt=True,
-    decode_kwargs=dict(skip_special_tokens=True)
-)
+# ======== 新增：生成与输入长度限制 / 加速开关 ========
+MAX_INPUT_TOKENS = 3000         # 评审输入截断上限（可 2500~3500 之间调）
+MAX_NEW_TOKENS   = 512          # 生成上限（必要时可到 640/768）
+FAST_SUMMARY_TOK = 480          # 超长文本“先摘要”的输出上限
+LONG_TEXT_THRESHOLD_TOK = 6000  # 超过此上下文 token 视为超长，触发“先摘要”
 
-# 允许 TF32（可选加速；不支持也不会报错）
+# 允许 TF32（若硬件支持可提速；不支持也不会报错）
 try:
     torch.backends.cuda.matmul.allow_tf32 = True
     torch.backends.cudnn.allow_tf32 = True
 except Exception:
     pass
+# ======== 新增结束 ========
 
-# =========================
-# 推理时长保护：输入/输出限长 + 超长先摘要
-# =========================
-MAX_INPUT_TOKENS = 3000          # 评审输入上限（按 token）
-MAX_NEW_TOKENS   = 512           # 生成上限（按 token）
-FAST_SUMMARY_TOK = 480           # 超长文本的“先摘要”输出上限
-LONG_TEXT_THRESHOLD_TOK = 6000   # 触发“先摘要”的 token 阈值
-
-def truncate_by_tokens(text: str, tokenizer, max_tokens: int) -> str:
-    """按 token 截断文本，优先保留头尾内容。"""
-    ids = tokenizer(text, return_tensors="pt", add_special_tokens=False)["input_ids"][0]
-    if ids.size(0) <= max_tokens:
-        return text
-    head = max_tokens // 2
-    tail = max_tokens - head
-    kept = torch.cat([ids[:head], ids[-tail:]], dim=0)
-    return tokenizer.decode(kept, skip_special_tokens=True)
-
-# =========================
-# 轻量“是否为论文”检测
-# =========================
+# Initialize AI classifier for paper detection
+# Using a lightweight model for text classification
 try:
     paper_classifier = pipeline(
         "text-classification",
-        model="fabriceyhc/bert-base-uncased-arxiv-classification",
+        model="fabriceyhc/bert-base-uncased-arxiv-classification",  # You can use other models
         device=0 if torch.cuda.is_available() else -1,
         truncation=True,
         max_length=512
@@ -95,10 +74,11 @@ except Exception as e:
     paper_classifier = None
     AI_CLASSIFIER_AVAILABLE = False
 
+# Alternative: Use zero-shot classification (more flexible but slower)
 def init_zero_shot_classifier():
     try:
-        from transformers import pipeline as hfpipe
-        classifier = hfpipe(
+        from transformers import pipeline
+        classifier = pipeline(
             "zero-shot-classification",
             model="facebook/bart-large-mnli",
             device=0 if torch.cuda.is_available() else -1
@@ -108,62 +88,117 @@ def init_zero_shot_classifier():
         print(f"Could not initialize zero-shot classifier: {e}")
         return None
 
-zero_shot_classifier = init_zero_shot_classifier() if not AI_CLASSIFIER_AVAILABLE else None
+# If primary classifier fails, try zero-shot
+if not AI_CLASSIFIER_AVAILABLE:
+    zero_shot_classifier = init_zero_shot_classifier()
+    if zero_shot_classifier:
+        AI_CLASSIFIER_AVAILABLE = True
+else:
+    zero_shot_classifier = None
 
 def ai_check_paper(text):
-    """AI 模型判断是否为研究论文：(is_paper, confidence, reason)"""
-    if not AI_CLASSIFIER_AVAILABLE and not zero_shot_classifier:
+    """
+    Use AI model to check if text is a research paper
+    Returns (is_paper, confidence, ai_reason)
+    """
+    if not AI_CLASSIFIER_AVAILABLE:
         return None, 0, "AI classifier not available"
-
+    
+    # Truncate text for AI model (keep beginning and end which are most informative)
     max_chars = 2000
     if len(text) > max_chars * 2:
         text_sample = text[:max_chars] + "\n...\n" + text[-max_chars:]
     else:
-        text_sample = text[:max_chars * 2]
-
+        text_sample = text[:max_chars*2]
+    
     try:
         if zero_shot_classifier and not paper_classifier:
+            # Use zero-shot classification
             labels = [
-                "academic research paper", "scientific article", "technical report",
-                "business document", "news article", "blog post", "other document"
+                "academic research paper",
+                "scientific article", 
+                "technical report",
+                "business document",
+                "news article",
+                "blog post",
+                "other document"
             ]
+            
             result = zero_shot_classifier(
                 text_sample,
                 candidate_labels=labels,
                 hypothesis_template="This text is a {}."
             )
-            top_label = result["labels"][0]
-            top_score = result["scores"][0]
+            
+            # Check if top labels are paper-related
+            top_label = result['labels'][0]
+            top_score = result['scores'][0]
+            
             paper_labels = {"academic research paper", "scientific article", "technical report"}
+            
             if top_label in paper_labels:
                 if top_score > 0.7:
-                    return True, top_score, f"AI detected: {top_label} ({top_score:.2f})"
+                    return True, top_score, f"AI detected: {top_label} (confidence: {top_score:.2f})"
                 elif top_score > 0.5:
-                    return True, 0.6, f"AI detected: likely {top_label} ({top_score:.2f})"
+                    return True, 0.6, f"AI detected: likely {top_label} (confidence: {top_score:.2f})"
                 else:
-                    return False, top_score, "AI uncertain"
+                    return False, top_score, f"AI detected: uncertain document type"
             else:
-                return False, 1 - top_score, f"AI: {top_label}, not a research paper"
+                return False, 1-top_score, f"AI detected: {top_label}, not a research paper"
+                
         elif paper_classifier:
+            # Use pre-trained classifier
             result = paper_classifier(text_sample)[0]
-            label = result["label"].lower()
-            score = result["score"]
+            
+            # Check if the label indicates it's a paper
+            # This depends on the specific model used
+            label = result['label'].lower()
+            score = result['score']
+            
+            # Adjust based on your chosen model's labels
             paper_keywords = ['cs', 'math', 'physics', 'eess', 'econ', 'stat', 'q-bio']
-            is_paper = any(k in label for k in paper_keywords)
+            is_paper = any(keyword in label for keyword in paper_keywords)
+            
             if is_paper:
-                return True, score, f"AI detected paper: {label} ({score:.2f})"
+                return True, score, f"AI detected: {label} paper (confidence: {score:.2f})"
             else:
-                return False, 1 - score, "AI: not a research paper"
+                return False, 1-score, f"AI detected: not a research paper"
+                
     except Exception as e:
         print(f"AI classification error: {e}")
         return None, 0, "AI classification failed"
+    
     return None, 0, "AI check not performed"
 
+# Define prompts
 SYSTEM_PROMPT_TEMPLATE = """You are an expert reviewer for AI conferences. You follow best practices and review papers according to the reviewer guidelines.
 
 Reviewer guidelines:
-1. Read the paper...
-[snip — unchanged for brevity in comments]
+1. Read the paper: It's important to carefully read through the entire paper, and to look up any related work and citations that will help you comprehensively evaluate it. Be sure to give yourself sufficient time for this step.
+2. While reading, consider the following:
+    - Objective of the work: What is the goal of the paper? Is it to better address a known application or problem, draw attention to a new application or problem, or to introduce and/or explain a new theoretical finding? A combination of these? Different objectives will require different considerations as to potential value and impact.
+    - Strong points: is the submission clear, technically correct, experimentally rigorous, reproducible, does it present novel findings (e.g., theoretically, algorithmically, etc.)?
+    - Weak points: is it weak in any of the aspects listed in b.?
+    - Be mindful of potential biases and try to be open-minded about the value and interest a paper can hold for the community, even if it may not be very interesting for you.
+3. Answer four key questions for yourself, to make a recommendation to Accept or Reject:
+    - What is the specific question and/or problem tackled by the paper?
+    - Is the approach well motivated, including being well-placed in the literature?
+    - Does the paper support the claims? This includes determining if results, whether theoretical or empirical, are correct and if they are scientifically rigorous.
+    - What is the significance of the work? Does it contribute new knowledge and sufficient value to the community? Note, this does not necessarily require state-of-the-art results. Submissions bring value to the community when they convincingly demonstrate new, relevant, impactful knowledge (incl., empirical, theoretical, for practitioners, etc).
+4. Write your review including the following information: 
+    - Summarize what the paper claims to contribute. Be positive and constructive.
+    - List strong and weak points of the paper. Be as comprehensive as possible.
+    - Clearly state your initial recommendation (accept or reject) with one or two key reasons for this choice.
+    - Provide supporting arguments for your recommendation.
+    - Ask questions you would like answered by the authors to help you clarify your understanding of the paper and provide the additional evidence you need to be confident in your assessment.
+    - Provide additional feedback with the aim to improve the paper. Make it clear that these points are here to help, and not necessarily part of your decision assessment.
+
+Your write reviews in markdown format. Your reviews contain the following sections:
+
+# Review
+
+{review_fields}
+
 Your response must only contain the review in markdown format with sections as defined above.
 """
 
@@ -172,18 +207,146 @@ USER_PROMPT_TEMPLATE = """Review the following paper:
 {paper_text}
 """
 
+# For now, use fixed review fields
 REVIEW_FIELDS = """## Summary
-[snip — unchanged template content]
+Briefly summarize the paper and its contributions. This is not the place to critique the paper; the authors should generally agree with a well-written summary.
+
+## Novelty
+Please assign the paper a numerical rating on the following scale to indicate the novelty and originality of the work. Consider whether the paper presents new ideas, methods, or perspectives that have not been explored before. Choose from the following:
+4: excellent - Highly original work with groundbreaking ideas or completely novel approaches
+3: good - Significant new contributions with clear advances over existing work
+2: fair - Some new elements but largely incremental improvements or combinations of existing ideas
+1: poor - Little to no novelty, mostly reproducing existing work or trivial variations
+
+## Novelty Explanation
+IMPORTANT: Focus ONLY on novelty aspects. DO NOT discuss soundness, presentation, or general contribution here.
+Please provide specific justification for your novelty score by addressing:
+- What specific new concepts, methods, or approaches does this paper introduce?
+- How do these differ from existing work in the field? Cite specific prior work for comparison.
+- Are the differences substantial or incremental?
+- Is this addressing a problem in a genuinely new way, or applying known methods to a new domain?
+DO NOT repeat content from other sections. DO NOT discuss writing quality, experimental rigor, or implementation details here.
+
+## Soundness
+Please assign the paper a numerical rating on the following scale to indicate the soundness of the technical claims, experimental and research methodology and on whether the central claims of the paper are adequately supported with evidence. Choose from the following:
+4: excellent
+3: good
+2: fair
+1: poor
+
+## Soundness Explanation
+IMPORTANT: Focus ONLY on technical correctness and methodological rigor. DO NOT discuss novelty or writing quality here.
+Please provide specific reasons for your soundness score by addressing:
+- Are the technical claims mathematically/logically correct?
+- Is the experimental methodology rigorous and appropriate?
+- Are the experiments sufficient to support the claims?
+- Are there any methodological flaws or missing controls?
+- Is the statistical analysis (if any) appropriate and correctly executed?
+DO NOT repeat content from other sections. DO NOT discuss the novelty of the approach or presentation quality here.
+
+## Presentation
+Please assign the paper a numerical rating on the following scale to indicate the quality of the presentation. This should take into account the writing style and clarity, as well as contextualization relative to prior work. Choose from the following:
+4: excellent
+3: good
+2: fair
+1: poor
+
+## Presentation Explanation
+IMPORTANT: Focus ONLY on writing quality, clarity, and organization. DO NOT discuss technical merit or novelty here.
+Please explain your presentation score by addressing:
+- Is the paper well-organized and easy to follow?
+- Are the main ideas clearly explained?
+- Are figures, tables, and visualizations effective and well-designed?
+- Is the related work section comprehensive and fair?
+- Are mathematical notations consistent and clear?
+- Is the language precise and grammatically correct?
+DO NOT repeat content from other sections. DO NOT discuss the novelty of ideas or soundness of methods here.
+
+## Contribution
+Please assign the paper a numerical rating on the following scale to indicate the quality of the overall contribution this paper makes to the research area being studied. Are the questions being asked important? Does the paper bring a significant originality of ideas and/or execution? Are the results valuable to share with the broader ICLR community? Choose from the following:
+4: excellent
+3: good
+2: fair
+1: poor
+
+## Contribution Explanation
+IMPORTANT: Focus on the OVERALL IMPACT and SIGNIFICANCE to the field. This is different from novelty.
+Please justify your contribution score by explaining:
+- Why is this work important for the field?
+- What practical or theoretical impact could this have?
+- Who would benefit from this work and how?
+- Does this open new research directions or close important gaps?
+- How significant are the improvements over baselines (if applicable)?
+Consider both immediate utility and long-term impact. DO NOT simply repeat the novelty assessment here.
+
+## Strengths
+List the main strengths of the paper. Be specific and provide evidence. Each strength should be a separate bullet point. Focus on what the paper does well across all dimensions (novelty, soundness, presentation, contribution). Avoid generic statements.
+
+## Weaknesses
+List the main weaknesses of the paper. Be specific, constructive, and actionable. Each weakness should be a separate bullet point with suggestions for improvement where possible. Focus on significant issues that affect the paper's validity or impact.
+
+## Questions
+List specific questions for the authors that could clarify ambiguities or address concerns. Number each question. These should be questions where the answer could potentially change your assessment of the paper.
+
+## Flag For Ethics Review
+If there are ethical issues with this paper, please flag the paper for an ethics review and select area of expertise that would be most useful for the ethics reviewer to have. Please select all that apply. Choose from the following:
+No ethics review needed.
+Yes, Discrimination / bias / fairness concerns
+Yes, Privacy, security and safety
+Yes, Legal compliance (e.g., GDPR, copyright, terms of use)
+Yes, Potentially harmful insights, methodologies and applications
+Yes, Responsible research practice (e.g., human subjects, data release)
+Yes, Research integrity issues (e.g., plagiarism, dual submission)
+Yes, Unprofessional behaviors (e.g., unprofessional exchange between authors and reviewers)
+Yes, Other reasons (please specify below)
+
+## Details Of Ethics Concerns
+Please provide details of your concerns. If no ethics review is needed, write "N/A".
+
+## Rating
+Please provide an "overall score" for this submission. Choose from the following:
+1: strong reject
+3: reject, not good enough
+5: marginally below the acceptance threshold
+6: marginally above the acceptance threshold
+8: accept, good paper
+10: strong accept, should be highlighted at the conference
+
 ## Overall Justification
-[snip]
+Provide a comprehensive justification for your overall rating that:
+- Synthesizes the assessments from all dimensions (novelty, soundness, presentation, contribution)
+- Explains how you weighted different aspects in arriving at your final score
+- Clearly states whether the strengths outweigh the weaknesses or vice versa
+- Indicates what would need to change for a different rating
+This should be a holistic assessment, not a repetition of individual sections.
+
 """
 
+# ======== 新增：token 级截断 ========
+def truncate_by_tokens(text: str, tokenizer, max_tokens: int) -> str:
+    """按 token 截断文本，优先保留开头+结尾（论文信息密度最高的区域）"""
+    ids = tokenizer(text, return_tensors="pt", add_special_tokens=False)["input_ids"][0]
+    if ids.size(0) <= max_tokens:
+        return text
+    head = max_tokens // 2
+    tail = max_tokens - head
+    kept = torch.cat([ids[:head], ids[-tail:]], dim=0)
+    return tokenizer.decode(kept, skip_special_tokens=True)
+# ======== 新增结束 ========
+
+# Enhanced function that combines rule-based and AI checks
 def is_research_paper(text, use_ai=True):
-    """规则 +（可选）AI 检测是否像论文。返回 (is_paper, confidence, reason)"""
+    """
+    Check if the given text appears to be a research paper.
+    Combines rule-based detection with AI classification.
+    Returns (is_paper, confidence, reason)
+    """
     if not text or len(text) < MIN_LENGTH:
         return False, 0, "Text is too short to be a research paper"
-
+    
     text_lower = text.lower()
+    
+    # Academic paper indicators (must have multiple)
     indicators = {
         'abstract': bool(re.search(r'\babstract\b', text_lower)),
         'introduction': bool(re.search(r'\bintroduction\b', text_lower)),
@@ -191,73 +354,114 @@ def is_research_paper(text, use_ai=True):
         'references': bool(re.search(r'\b(references|bibliography)\b', text_lower)),
         'methodology': bool(re.search(r'\b(method|methodology|approach|algorithm|model)\b', text_lower)),
         'results': bool(re.search(r'\b(results|experiments|evaluation|analysis)\b', text_lower)),
-        'citations': bool(re.search(r'\[[\d,\s]+\]|\(\w+,?\s*\d{4}\)', text)),
+        'citations': bool(re.search(r'\[[\d,\s]+\]|\(\w+,?\s*\d{4}\)', text)),  # [1,2,3] or (Author, 2024)
         'figures_tables': bool(re.search(r'\b(figure\s*\d+|table\s*\d+|fig\.\s*\d+)\b', text_lower)),
         'academic_terms': bool(re.search(r'\b(propose|present|demonstrate|evaluate|contribution|novel|state-of-the-art)\b', text_lower))
     }
+    
+    # Count how many indicators are present
     indicator_count = sum(indicators.values())
-
+    
+    # Check for non-paper content with context
+    # Only flag as non-paper if these terms appear WITHOUT academic context
     non_paper_indicators = []
+    
+    # Check for invoice/receipt patterns (multiple commercial terms together)
     if re.search(r'\b(invoice|receipt)\b', text_lower) and re.search(r'\b(total|amount|payment|billing)\b', text_lower):
         non_paper_indicators.append(True)
+    
+    # Check for purchase order specifically (not just "purchase")
     if re.search(r'\bpurchase order\b', text_lower):
         non_paper_indicators.append(True)
+    
+    # Check for letter format
     if re.search(r'\b(dear\s+(sir|madam|customer)|sincerely|best regards|yours truly)\b', text_lower):
         non_paper_indicators.append(True)
+    
+    # Check for textbook structure
     if re.search(r'\b(chapter\s+\d+|lesson\s+\d+|exercise\s+\d+)\b', text_lower) and indicator_count < 3:
         non_paper_indicators.append(True)
+    
+    # Check for HTML/web content
     if re.search(r'<html|<body|<div|<script|<!DOCTYPE', text_lower):
         non_paper_indicators.append(True)
+    
+    # Check for recipe/cooking content
     if re.search(r'\b(ingredients|recipe|preparation|cooking time|servings)\b', text_lower) and not re.search(r'\b(algorithm|method|experiment)\b', text_lower):
         non_paper_indicators.append(True)
+    
+    # If we have strong non-paper indicators AND weak academic indicators, it's not a paper
     if any(non_paper_indicators) and indicator_count < 6:
         return False, 0, "Content appears to be a non-academic document"
-
+    
+    # Get AI assessment if available and requested
     ai_result = None
     ai_confidence = 0
     ai_reason = ""
-    if use_ai and (AI_CLASSIFIER_AVAILABLE or zero_shot_classifier):
+    
+    if use_ai and AI_CLASSIFIER_AVAILABLE:
         ai_result, ai_confidence, ai_reason = ai_check_paper(text)
-
+    
+    # Combine rule-based and AI assessments
+    # Rule-based decision logic
     if indicator_count == 9:
-        rule_decision, rule_confidence = True, 0.9
-        rule_reason = f"Found all {indicator_count}/9 academic indicators"
+        rule_decision = True
+        rule_confidence = 0.9
+        rule_reason = f"Found all {indicator_count}/9 academic paper indicators"
     elif indicator_count >= 6:
-        rule_decision, rule_confidence = True, 0.6
+        rule_decision = True
+        rule_confidence = 0.6
         missing = [k for k, v in indicators.items() if not v]
-        rule_reason = f"Found {indicator_count}/9 indicators. Missing: {', '.join(missing)}"
+        rule_reason = f"Found only {indicator_count}/9 indicators. Missing: {', '.join(missing)}"
     else:
-        rule_decision, rule_confidence = False, 0
+        rule_decision = False
+        rule_confidence = 0
         missing = [k for k, v in indicators.items() if not v]
         rule_reason = f"Found only {indicator_count}/9 indicators. Missing: {', '.join(missing[:4])}"
-
+    
+    # Combine decisions
     if ai_result is not None:
+        # Both methods available - combine them
+        # Weight: 60% rule-based, 40% AI
         combined_confidence = (rule_confidence * 0.6) + (ai_confidence * 0.4)
+        
+        # Decision logic
         if rule_decision and ai_result:
-            return True, (0.9 if combined_confidence >= 0.9 else 0.6), f"{rule_reason}. {ai_reason}"
+            # Both agree it's a paper
+            if combined_confidence >= 0.9:
+                return True, 0.9, f"High confidence: {rule_reason}. {ai_reason}"
+            else:
+                return True, 0.6, f"Warning: {rule_reason}. {ai_reason}"
         elif not rule_decision and not ai_result:
+            # Both agree it's not a paper
             return False, 0, f"Not a research paper. {rule_reason}. {ai_reason}"
         else:
+            # Disagreement - use weighted decision
             if combined_confidence >= 0.5:
                 return True, 0.6, f"Mixed signals: {rule_reason}. {ai_reason}"
             else:
                 return False, 0, f"Likely not a research paper. {rule_reason}. {ai_reason}"
     else:
+        # Only rule-based available
         if rule_decision:
-            return True, (0.9 if rule_confidence >= 0.9 else 0.6), rule_reason
+            if rule_confidence >= 0.9:
+                return True, 0.9, f"High confidence: {rule_reason}"
+            else:
+                return True, 0.6, f"Warning: {rule_reason}"
         else:
             return False, 0, f"Does not appear to be a research paper. {rule_reason}"
 
 def create_messages(review_fields, paper_text):
-    return [
+    messages = [
         {"role": "system", "content": SYSTEM_PROMPT_TEMPLATE.format(review_fields=review_fields)},
-        {"role": "user",   "content": USER_PROMPT_TEMPLATE.format(paper_text=paper_text)},
+        {"role": "user", "content": USER_PROMPT_TEMPLATE.format(paper_text=paper_text)},
     ]
+    return messages
 
 @spaces.GPU()
 def convert_file(filepath):
     full_text, images, out_metadata = convert_single_pdf(
-        filepath, model_refs, metadata=metadata, max_pages=MAX_PAGES
+            filepath, model_refs, metadata=metadata, max_pages=MAX_PAGES
     )
     return full_text
 
@@ -269,10 +473,12 @@ def process_file(file):
             filetype = find_filetype(filepath)
             if filetype == "other":
                 raise ValueError()
+
             length = get_length_of_text(filepath)
             if length < MIN_LENGTH:
                 raise ValueError()
-        paper_text = convert_file(filepath).strip()
+        paper_text = convert_file(filepath)
+        paper_text = paper_text.strip()
         if not len(paper_text) > MIN_LENGTH:
             raise ValueError()
     except spaces.zero.gradio.HTMLError as e:
@@ -282,33 +488,27 @@ def process_file(file):
         print(traceback.format_exc())
         print(f"Error converting {filepath}: {e}")
         return "Error processing pdf", False
-
+    
+    # Check if it's a research paper (with AI)
     is_paper, confidence, reason = is_research_paper(paper_text, use_ai=True)
     if not is_paper:
-        return (
-            "⚠️ **Not a Research Paper**\n\n"
-            "The uploaded document does not appear to be a research paper.\n\n"
-            f"Reason: {reason}\n\n"
-            "Please upload a proper academic/research paper with sections like Abstract, Introduction, Methodology, Results, and References.",
-            False,
-        )
+        return f"⚠️ **Not a Research Paper**\n\nThe uploaded document does not appear to be a research paper.\n\nReason: {reason}\n\nPlease upload a proper academic/research paper with sections like Abstract, Introduction, Methodology, Results, and References.", False
+    
+    # If confidence is low (6-8 indicators), add a warning
     if confidence < 0.9:
-        paper_text = (
-            f"⚠️ **Warning**: {reason}. \n\n"
-            "The document may be incomplete or missing key sections. Proceeding with review generation...\n\n---\n\n"
-            f"{paper_text}"
-        )
+        paper_text = f"⚠️ **Warning**: {reason}. \n\nThe document may be incomplete or missing key sections. Proceeding with review generation...\n\n---\n\n{paper_text}"
+    
     return paper_text, True
 
-# ============== 新增：一步式端点（自动瘦身） ==============
+# ======== 新增：一步式评审端点（自动“先摘要后评审” + 截断） ========
 @spaces.GPU(duration=90)
 def review(paper_text: str, review_template: str):
-    # 快速检查（禁用 AI，加速）
+    # Quick paper check（不用 AI 分类，加速）
     is_paper, confidence, reason = is_research_paper(paper_text, use_ai=False)
     if not is_paper:
         return f"⚠️ Cannot generate review: {reason}"
 
-    # 估算 token，超长则先摘要
+    # 检测是否超长，必要时先摘要
     pt_ids = tokenizer(paper_text, return_tensors="pt", add_special_tokens=False)["input_ids"][0]
     if pt_ids.size(0) > LONG_TEXT_THRESHOLD_TOK:
         summary_prompt = (
@@ -316,7 +516,7 @@ def review(paper_text: str, review_template: str):
         )
         messages = [
             {"role": "system", "content": "You are a helpful scientific summarizer."},
-            {"role": "user",   "content": f"{summary_prompt}\n\n---\n\n{paper_text}"},
+            {"role": "user", "content": f"{summary_prompt}\n\n---\n\n{paper_text}"}
         ]
         sum_ids = tokenizer.apply_chat_template(
             messages, add_generation_prompt=True, return_tensors="pt"
@@ -329,17 +529,17 @@ def review(paper_text: str, review_template: str):
         )
         paper_text = tokenizer.decode(sum_out[0][sum_ids.shape[1]:], skip_special_tokens=True)
 
-    # 截断评审输入，确保不超限
+    # 截断评审输入，保证不会过长
     paper_text = truncate_by_tokens(paper_text, tokenizer, MAX_INPUT_TOKENS)
 
-    # 构造消息并生成（非流式，一步返回）
+    # 构建评审消息并一次性生成（非流式，适合 API 调用）
     messages = create_messages(review_template, paper_text)
     input_ids = tokenizer.apply_chat_template(
         messages, add_generation_prompt=True, return_tensors='pt'
     ).to(model.device)
-    attention_mask = torch.ones_like(input_ids)
 
-    out = model.generate(
+    attention_mask = torch.ones_like(input_ids)
+    generation_kwargs = dict(
         input_ids=input_ids,
         attention_mask=attention_mask,
         max_new_tokens=MAX_NEW_TOKENS,
@@ -348,34 +548,38 @@ def review(paper_text: str, review_template: str):
         top_p=0.9,
         pad_token_id=tokenizer.pad_token_id
     )
+    out = model.generate(**generation_kwargs)
     return tokenizer.decode(out[0][input_ids.shape[1]:], skip_special_tokens=True).replace("<|eot_id|>", "")
+# ======== 新增结束 ========
 
-# ============== 改造：原 /generate（流式，但限长） ==============
 @spaces.GPU(duration=90)
 def generate(paper_text, review_template):
-    is_paper, confidence, reason = is_research_paper(paper_text, use_ai=False)
+    # Final check before generation
+    is_paper, confidence, reason = is_research_paper(paper_text, use_ai=False)  # Quick check without AI
     if not is_paper:
         return f"⚠️ Cannot generate review: {reason}"
 
-    # 关键：截断输入，防止上下文超大
+    # ——关键：先按 token 截断输入，避免上下文超大——
     paper_text = truncate_by_tokens(paper_text, tokenizer, MAX_INPUT_TOKENS)
-
+    
     messages = create_messages(review_template, paper_text)
     input_ids = tokenizer.apply_chat_template(
         messages,
         add_generation_prompt=True,
         return_tensors='pt'
     ).to(model.device)
+    
+    # Create attention mask
     attention_mask = torch.ones_like(input_ids)
-
+    
     print(f"input_ids shape (truncated): {input_ids.shape}")
     generation_kwargs = dict(
         input_ids=input_ids,
         attention_mask=attention_mask,
         streamer=streamer,
-        max_new_tokens=MAX_NEW_TOKENS,   # 由 4096 降为 512
+        max_new_tokens=MAX_NEW_TOKENS,   # from 4096 -> 512
         do_sample=True,
-        temperature=0.5,
+        temperature=0.5,                 # 略降温度，更稳更短
         top_p=0.9,
         pad_token_id=tokenizer.pad_token_id
     )
@@ -386,10 +590,9 @@ def generate(paper_text, review_template):
         generated_text += new_text
         yield generated_text.replace("<|eot_id|>", "")
 
-# =========================
-# UI（保持不变，仅加 api_name）
-# =========================
-title_html = """<h1 align="center">OpenReviewer</h1>
+# UI code remains exactly the same...
+# ui
+title = """<h1 align="center">OpenReviewer</h1>
 <div align="center">Using <a href="https://huggingface.co/maxidl/Llama-OpenReviewer-8B" target="_blank"><code>Llama-OpenReviewer-8B</code></a> - Built with Llama</div>
 """
 
@@ -397,36 +600,45 @@ description = """This is an online demo featuring [Llama-OpenReviewer-8B](https:
 
 ## Demo Guidelines
 
-1. Upload your paper as a PDF file. Alternatively you can paste the full text of your paper in markdown format below...
+1. Upload your paper as a PDF file. Alternatively you can paste the full text of your paper in markdown format below. We do **not** store your data. User data is kept in ephemeral storage during processing.
+
+2. Once you upload a PDF, it will be converted to markdown and **validated to ensure it's a research paper**. This takes some time as it runs multiple transformer models to parse the layout and extract text and tables. Check out [marker](https://github.com/VikParuchuri/marker/tree/master) for details.
+
+3. Having obtained a markdown version of your paper and confirmed it's a valid research paper, you can now click *Generate Review*.
+
+Take a look at the Review Template to properly interpret the generated review. You can also change the review template before generating in case you want to generate a review with a different schema and aspects.
+
+To obtain more than one review, just generate again.
+
+**GPU quota:** If exceeded, either sign in with your HF account or come back later. Your quota has a half-life of 2 hours.
+
 """
 
 theme = gr.themes.Default(primary_hue="gray", secondary_hue="blue", neutral_hue="slate")
 with gr.Blocks(theme=theme) as demo:
-    gr.HTML(title_html)
-    gr.Markdown(description)
-
+    title = gr.HTML(title)
+    description = gr.Markdown(description)
+    
+    # Add paper validation status
     with gr.Row():
         file_input = gr.File(file_types=[".pdf"], file_count="single")
         validation_status = gr.Markdown("", visible=False)
-
-    paper_text_field = gr.Textbox(
-        "Upload a pdf or paste the full text of your paper in markdown format here.",
-        label="Paper Text",
-        lines=20, max_lines=20, autoscroll=False
-    )
-
+    
+    paper_text_field = gr.Textbox("Upload a pdf or paste the full text of your paper in markdown format here.", label="Paper Text", lines=20, max_lines=20, autoscroll=False)
+    
     with gr.Accordion("Review Template", open=False):
-        gr.Markdown("We use the ICLR 2025 review template by default, but you can modify the template below as you like.")
-        review_template_field = gr.Textbox(label=" ", lines=20, max_lines=20, autoscroll=False, value=REVIEW_FIELDS)
-
+        review_template_description = gr.Markdown("We use the ICLR 2025 review template by default, but you can modify the template below as you like.")
+        review_template_field = gr.Textbox(label=" ",lines=20, max_lines=20, autoscroll=False, value=REVIEW_FIELDS)
+    
     generate_button = gr.Button("Generate Review", interactive=False)
-
+    
     def handle_file_upload(file):
         if file is None:
             return "", gr.update(visible=False), gr.update(interactive=False)
         text, is_valid = process_file(file)
         if is_valid:
-            is_paper, confidence, reason = is_research_paper(text, use_ai=False)
+            # Check confidence level for appropriate message
+            is_paper, confidence, reason = is_research_paper(text, use_ai=False)  # Quick check for display
             if confidence >= 0.9:
                 status_msg = "✅ **Document validated**: This appears to be a complete research paper."
             else:
@@ -434,48 +646,49 @@ with gr.Blocks(theme=theme) as demo:
             return text, gr.update(value=status_msg, visible=True), gr.update(interactive=True)
         else:
             return text, gr.update(value="❌ **Validation failed**: Please upload a research paper.", visible=True), gr.update(interactive=False)
-
+    
     def handle_text_change(text):
         if not text or len(text) < 200:
             return gr.update(interactive=False), gr.update(visible=False)
+        
         is_paper, confidence, reason = is_research_paper(text, use_ai=True)
         if is_paper:
-            status = ("✅ **Text validated**: This appears to be a complete research paper."
-                      if confidence >= 0.9 else
-                      f"⚠️ **Warning**: {reason}\n\nThe document may be incomplete or missing key sections.")
+            if confidence >= 0.9:
+                status = "✅ **Text validated**: This appears to be a complete research paper."
+            else:
+                # confidence < 0.9 means warning (6-8 indicators)
+                status = f"⚠️ **Warning**: {reason}\n\nThe document may be incomplete or missing key sections."
             return gr.update(interactive=True), gr.update(value=status, visible=True)
         else:
             return gr.update(interactive=False), gr.update(value=f"❌ **Not a research paper**: {reason}", visible=True)
-
+    
     file_input.upload(handle_file_upload, file_input, [paper_text_field, validation_status, generate_button])
     paper_text_field.change(handle_text_change, paper_text_field, [generate_button, validation_status])
-
+    
     review_field = gr.Markdown("\n\n\n\n\n", label="Review")
 
-    # 显式命名 /generate（流式）
+    # 生成按钮（流式），显式命名 /generate
     generate_button.click(
-        fn=lambda: gr.update(interactive=False),
-        inputs=None,
+        fn=lambda: gr.update(interactive=False), 
+        inputs=None, 
         outputs=generate_button
     ).then(
-        generate,
-        [paper_text_field, review_template_field],
+        generate, 
+        [paper_text_field, review_template_field], 
         review_field,
-        api_name="generate"
+        api_name="generate"       # 显式命名
     ).then(
-        fn=lambda: gr.update(interactive=True),
-        inputs=None,
+        fn=lambda: gr.update(interactive=True), 
+        inputs=None, 
         outputs=generate_button
     )
 
-    # （隐藏）显式暴露“一步式” /review 端点，供 SDK 使用
-    gr.Button(visible=False).click(
-        review,
-        [paper_text_field, review_template_field],
-        review_field,
-        api_name="review"
+    # （可选）隐藏按钮以暴露“一步式” /review 端点，方便 SDK 直接调用
+    dummy = gr.Button(visible=False)
+    dummy.click(
+        review, [paper_text_field, review_template_field], review_field, api_name="review"
     )
-
+    
     demo.title = "OpenReviewer"
 
 if __name__ == "__main__":