Update app.py

app.py

@@ -1,64 +1,245 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+import os
+import spaces
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import PyPDF2
+from io import BytesIO
+import torch
 
+# 设置环境变量
+HF_TOKEN = os.environ.get("HF_TOKEN", None)
+
+DESCRIPTION = '''
+<div>
+<h1 style="text-align: center;">Academic Paper Improver</h1>
+<p>This Space helps you improve sections of your academic paper using the <a href="https://huggingface.co/Xtra-Computing/XtraGPT-7B"><b>XtraGPT-7B</b></a> model.</p>
+<p>Upload your PDF paper, select a section of text you want to improve, and specify your requirements.</p>
+</div>
+'''
+
+CITATION = """
+<div style="font-family: monospace; white-space: pre; margin-top: 20px; line-height: 1.2;">
+@misc{XtraGPT,
+    title = {XtraGPT},
+    url = {https://huggingface.co/Xtra-Computing/XtraGPT-7B},
+    author = {Nuo Chen, Andre Lin HuiKai, Junyi Hou, Zining Zhang, Qian Wang, Xidong Wang, Bingsheng He},
+    month = {March},
+    year = {2025}
+}
+</div>
 """
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
+
+LICENSE = """
+<p/>
+---
+Built with XtraGPT-7B
 """
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
+css = """
+h1 {
+  text-align: center;
+  display: block;
+}
+#duplicate-button {
+  margin: auto;
+  color: white;
+  background: #1565c0;
+  border-radius: 100vh;
+}
+"""
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
+# 默认论文内容
+default_paper_content = """
+The dominant sequence transduction models are based on complex recurrent or convolutional neural networks in an encoder-decoder configuration. 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. Our model achieves 28.4 BLEU on the WMT 2014 English-to-German translation task, improving over the existing best results, including ensembles by over 2 BLEU. On the WMT 2014 English-to-French translation task, our model establishes a new single-model state-of-the-art BLEU score of 41.8 after training for 3.5 days on eight GPUs, a small fraction of the training costs of the best models from the literature. We show that the Transformer generalizes well to other tasks by applying it successfully to English constituency parsing both with large and limited training data.
+"""
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+# 直接加载模型和分词器
+tokenizer = AutoTokenizer.from_pretrained("Xtra-Computing/XtraGPT-7B")
+model = AutoModelForCausalLM.from_pretrained("Xtra-Computing/XtraGPT-7B", device_map="auto")
 
-    messages.append({"role": "user", "content": message})
+def extract_text_from_pdf(pdf_bytes):
+    """从上传的PDF文件中提取文本"""
+    if pdf_bytes is None:
+        return default_paper_content
+    
+    try:
+        # 确保pdf_bytes是字节类型
+        if isinstance(pdf_bytes, str):
+            return pdf_bytes  # 如果已经是字符串，直接返回
+            
+        # 直接使用字节对象
+        pdf_reader = PyPDF2.PdfReader(BytesIO(pdf_bytes))
+        
+        # 从所有页面提取文本
+        text = ""
+        for page_num in range(len(pdf_reader.pages)):
+            page = pdf_reader.pages[page_num]
+            text += page.extract_text() + "\n\n"
+        
+        # 限制文本长度，防止超出模型最大长度
+        if len(text) > 10000:  # 保守估计，留出足够空间给提示和生成
+            text = text[:10000] + "...(文本已截断)"
+            
+        return text
+    except Exception as e:
+        print(f"PDF提取错误: {str(e)}")
+        return default_paper_content
 
-    response = ""
+@spaces.GPU(duration=120)
+def improve_paper_section(paper_content, selected_content, improvement_prompt, temperature=0.1, max_new_tokens=512):
+    """
+    改进学术论文的一个部分 - 使用非流式生成
+    """
+    # 检查输入
+    if not selected_content or not improvement_prompt:
+        return "请同时提供要改进的文本和改进要求。"
+    
+    try:
+        # 限制paper_content长度，防止超出模型最大长度
+        if len(paper_content) > 10000:  # 保守估计
+            paper_content = paper_content[:10000] + "...(文本已截断)"
+        
+        # 构建提示
+        content = f"""
+Please improve the selected content based on the following. Act as an expert model for improving articles **PAPER_CONTENT**.
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
+The output needs to answer the **QUESTION** on **SELECTED_CONTENT** in the input. Avoid adding unnecessary length, unrelated details, overclaims, or vague statements.
+Focus on clear, concise, and evidence-based improvements that align with the overall context of the paper.
 
-        response += token
-        yield response
+<PAPER_CONTENT>
+{paper_content}
+</PAPER_CONTENT>
 
+<SELECTED_CONTENT>
+{selected_content}
+</SELECTED_CONTENT>
 
+<QUESTION>
+{improvement_prompt}
+</QUESTION>
 """
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
+        
+        # 准备输入
+        messages = [
+            {"role": "user", "content": content}
+        ]
+        
+        text = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+        
+        # 检查输入长度并截断
+        input_tokens = tokenizer.encode(text)
+        if len(input_tokens) > 15000:  # 为生成留出空间
+            input_tokens = input_tokens[:15000]
+            text = tokenizer.decode(input_tokens)
+            print(f"输入已截断至15000个token")
+        
+        # 使用非流式方式生成
+        input_ids = tokenizer.encode(text, return_tensors="pt").to(model.device)
+        
+        with torch.no_grad():
+            output_ids = model.generate(
+                input_ids,
+                max_new_tokens=max_new_tokens,
+                do_sample=(temperature > 0),
+                temperature=temperature if temperature > 0 else 1.0,
+                pad_token_id=tokenizer.eos_token_id
+            )
+        
+        # 只保留新生成的部分
+        generated_ids = output_ids[0, len(input_ids[0]):]
+        response = tokenizer.decode(generated_ids, skip_special_tokens=True)
+        
+        return response
+        
+    except Exception as e:
+        import traceback
+        error_details = traceback.format_exc()
+        print(f"生成错误: {str(e)}\n{error_details}")
+        return f"生成文本时出错: {str(e)}\n\n请尝试使用不同的参数或输入。"
 
+# 创建Gradio界面
+with gr.Blocks(fill_height=True, css=css) as demo:
+    # 存储提取的PDF文本
+    extracted_pdf_text = gr.State(default_paper_content)
+    
+    gr.Markdown(DESCRIPTION)
+    # gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button")
+    
+    with gr.Row():
+        with gr.Column():
+            # 步骤1：上传PDF
+            with gr.Group():
+                gr.Markdown("### Step 1: Upload your academic paper")
+                pdf_file = gr.File(
+                    label="Upload PDF",
+                    file_types=[".pdf"],
+                    type="binary"  # 直接获取二进制数据
+                )
+            
+            # 步骤2：提取并选择文本
+            with gr.Group():
+                gr.Markdown("### Step 2: Enter the text section to improve")
+                selected_content = gr.Textbox(
+                    label="Text to improve",
+                    placeholder="Paste the section of text you want to improve...",
+                    lines=5,
+                    value="The dominant sequence transduction models are based on complex recurrent or convolutional neural networks in an encoder-decoder configuration."
+                )
+            
+            # 步骤3：指定改进要求
+            with gr.Group():
+                gr.Markdown("### Step 3: Specify your improvement requirements")
+                improvement_prompt = gr.Textbox(
+                    label="Improvement requirements",
+                    placeholder="e.g., 'Make this more concise', 'Add more technical details', 'Redefine this concept'...",
+                    lines=3,
+                    value="help me make it more concise."
+                )
+            
+            with gr.Accordion("⚙️ Parameters", open=False):
+                temperature = gr.Slider(minimum=0, maximum=1, step=0.1, value=0.1, label="Temperature")
+                max_tokens = gr.Slider(minimum=128, maximum=1024, step=32, value=512, label="Max Tokens")
+            
+            submit_btn = gr.Button("Improve Text")
+        
+        with gr.Column():
+            # 输出
+            output = gr.Textbox(label="Improved Text", lines=20)
+            
+            # 显示提取的PDF文本（可折叠）
+            with gr.Accordion("Extracted PDF Content (for reference)", open=False):
+                pdf_content_display = gr.Textbox(
+                    label="Paper Content", 
+                    lines=10,
+                    value=default_paper_content
+                )
+    
+    # 当PDF上传时自动提取文本
+    def update_pdf_content(pdf_bytes):
+        if pdf_bytes is not None:
+            content = extract_text_from_pdf(pdf_bytes)
+            return content, content
+        return default_paper_content, default_paper_content
+    
+    pdf_file.change(
+        fn=update_pdf_content,
+        inputs=[pdf_file],
+        outputs=[extracted_pdf_text, pdf_content_display]
+    )
+    
+    # 处理文本改进
+    submit_btn.click(
+        fn=improve_paper_section,
+        inputs=[extracted_pdf_text, selected_content, improvement_prompt, temperature, max_tokens],
+        outputs=[output]
+    )
+    
+    # gr.Markdown(LICENSE)
+    gr.Markdown(CITATION)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()