2

app.py

@@ -1,3 +1,8 @@
+import subprocess
+
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"},
+               shell=True)
+
 import gradio as gr
 import torch
 from transformers import (
@@ -30,7 +35,10 @@ def convert_to_txt(file):
 model_name = "THUDM/LongCite-glm4-9b"
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
-    model_name, torch_dtype=torch.bfloat16, trust_remote_code=True, device_map="auto"
+    model_name,
+    torch_dtype=torch.bfloat16,
+    trust_remote_code=True,
+    attn_implementation="flash_attention_2",
 )
 
 html_styles = """<style>
@@ -110,20 +118,20 @@ def convert_to_html(statements, clicked=-1):
         html += "\n"
         if clicked == i:
             clicked_cite_html = (
-                html_styles
-                + """<br><span class="label">Citations of current statement:</span><br><div style="overflow-y: auto; padding: 20px; border: 0px dashed black; border-radius: 6px; background-color: #EFF2F6;">{}</div>""".format(
-                    "<br><br>\n".join(cite_html)
-                )
+                    html_styles
+                    + """<br><span class="label">Citations of current statement:</span><br><div style="overflow-y: auto; padding: 20px; border: 0px dashed black; border-radius: 6px; background-color: #EFF2F6;">{}</div>""".format(
+                "<br><br>\n".join(cite_html)
             )
-    all_cite_html = (
-        html_styles
-        + """<br><span class="label">All citations:</span><br>\n<div style="overflow-y: auto; padding: 20px; border: 0px dashed black; border-radius: 6px; background-color: #EFF2F6;">{}</div>""".format(
-            "<br><br>\n".join(all_cite_html).replace(
-                '<span class="highlight">', "<span>"
             )
-            if len(all_cite_html)
-            else "No citation in the answer"
+    all_cite_html = (
+            html_styles
+            + """<br><span class="label">All citations:</span><br>\n<div style="overflow-y: auto; padding: 20px; border: 0px dashed black; border-radius: 6px; background-color: #EFF2F6;">{}</div>""".format(
+        "<br><br>\n".join(all_cite_html).replace(
+            '<span class="highlight">', "<span>"
         )
+        if len(all_cite_html)
+        else "No citation in the answer"
+    )
     )
     return html, all_cite_html, clicked_cite_html, cite_num2idx
 
@@ -136,7 +144,7 @@ def render_context(file):
         raise gr.Error(f"ERROR: no uploaded document")
 
 
-@spaces.GPU()
+@spaces.GPU(duration=120)
 def run_llm(context, query):
     if not context:
         raise gr.Error("Error: no uploaded document")