init

.idea/.gitignore

@@ -0,0 +1,8 @@
+# Default ignored files
+/shelf/
+/workspace.xml
+# Editor-based HTTP Client requests
+/httpRequests/
+# Datasource local storage ignored files
+/dataSources/
+/dataSources.local.xml

.idea/LongWriter.iml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="PYTHON_MODULE" version="4">
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$" />
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+</module>

.idea/inspectionProfiles/Project_Default.xml

@@ -0,0 +1,21 @@
+<component name="InspectionProjectProfileManager">
+  <profile version="1.0">
+    <option name="myName" value="Project Default" />
+    <inspection_tool class="PyPackageRequirementsInspection" enabled="true" level="WARNING" enabled_by_default="true">
+      <option name="ignoredPackages">
+        <value>
+          <list size="8">
+            <item index="0" class="java.lang.String" itemvalue="openai" />
+            <item index="1" class="java.lang.String" itemvalue="sse_starlette" />
+            <item index="2" class="java.lang.String" itemvalue="fastapi" />
+            <item index="3" class="java.lang.String" itemvalue="timm" />
+            <item index="4" class="java.lang.String" itemvalue="gradio" />
+            <item index="5" class="java.lang.String" itemvalue="uvicorn" />
+            <item index="6" class="java.lang.String" itemvalue="diffusers" />
+            <item index="7" class="java.lang.String" itemvalue="transformers" />
+          </list>
+        </value>
+      </option>
+    </inspection_tool>
+  </profile>
+</component>

.idea/inspectionProfiles/profiles_settings.xml

@@ -0,0 +1,6 @@
+<component name="InspectionProjectProfileManager">
+  <settings>
+    <option name="USE_PROJECT_PROFILE" value="false" />
+    <version value="1.0" />
+  </settings>
+</component>

.idea/modules.xml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/LongWriter.iml" filepath="$PROJECT_DIR$/.idea/LongWriter.iml" />
+    </modules>
+  </component>
+</project>

.idea/vcs.xml

@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="VcsDirectoryMappings">
+    <mapping directory="" vcs="Git" />
+  </component>
+</project>

README.md

@@ -4,9 +4,14 @@ emoji: 💬
 colorFrom: yellow
 colorTo: purple
 sdk: gradio
-sdk_version: 4.36.1
+sdk_version: 4.41.1
+app_port: 7860
 app_file: app.py
 pinned: false
 ---
 
-An example chatbot using [Gradio](https://gradio.app), [`huggingface_hub`](https://huggingface.co/docs/huggingface_hub/v0.22.2/en/index), and the [Hugging Face Inference API](https://huggingface.co/docs/api-inference/index).
+# LongWriter
+
+```shell
+python app.py
+```

app.py

@@ -1,63 +1,116 @@
+from threading import Thread
+import spaces
 import gradio as gr
-from huggingface_hub import InferenceClient
-
-"""
-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
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
-
-
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    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
-
-        response += token
-        yield response
-
-"""
-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)",
-        ),
-    ],
+import torch
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    StoppingCriteria,
+    StoppingCriteriaList,
+    TextIteratorStreamer
 )
 
+model = AutoModelForCausalLM.from_pretrained("THUDM/longwriter-glm4-9b", trust_remote_code=True, device_map='auto')
+tokenizer = AutoTokenizer.from_pretrained("THUDM/longwriter-glm4-9b", trust_remote_code=True)
 
-if __name__ == "__main__":
-    demo.launch()
+
+class StopOnTokens(StoppingCriteria):
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        stop_ids = model.config.eos_token_id
+        for stop_id in stop_ids:
+            if input_ids[0][-1] == stop_id:
+                return True
+        return False
+
+
+@spaces.GPU()
+def predict(history, prompt, max_length, top_p, temperature):
+    stop = StopOnTokens()
+    messages = []
+    if prompt:
+        messages.append({"role": "system", "content": prompt})
+    for idx, (user_msg, model_msg) in enumerate(history):
+        if prompt and idx == 0:
+            continue
+        if idx == len(history) - 1 and not model_msg:
+            query = user_msg
+            break
+        if user_msg:
+            messages.append({"role": "user", "content": user_msg})
+        if model_msg:
+            messages.append({"role": "assistant", "content": model_msg})
+
+    model_inputs = tokenizer.build_chat_input(query, history=messages, role='user').input_ids.to(
+        next(model.parameters()).device)
+    streamer = TextIteratorStreamer(tokenizer, timeout=600, skip_prompt=True, skip_special_tokens=True)
+    eos_token_id = [tokenizer.eos_token_id, tokenizer.get_command("<|user|>"),
+                    tokenizer.get_command("<|observation|>")]
+    generate_kwargs = {
+        "input_ids": model_inputs,
+        "streamer": streamer,
+        "max_new_tokens": max_length,
+        "do_sample": True,
+        "top_p": top_p,
+        "temperature": temperature,
+        "stopping_criteria": StoppingCriteriaList([stop]),
+        "repetition_penalty": 1,
+        "eos_token_id": eos_token_id,
+    }
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
+    for new_token in streamer:
+        if new_token == '<|user|>':
+            continue
+        elif new_token:
+            history[-1][1] += new_token
+        yield history
+
+
+with gr.Blocks() as demo:
+    gr.Markdown(
+        """
+        <div style="text-align: center; font-size: 32px; font-weight: bold; margin-bottom: 20px;">
+            longwriter-glm4-9b Huggingface Space🤗
+        </div>
+        <div style="text-align: center;">
+            <a href="https://huggingface.co/THUDM/LongWriter-glm4-9b">🤗 Model Hub</a> |
+            <a href="https://github.com/THUDM/LongWriter">🌐 Github</a> |
+            <a href="https://arxiv.org/pdf/2408.07055">📜 arxiv </a>
+        </div>
+        """
+    )
+    chatbot = gr.Chatbot()
+
+    with gr.Row():
+        with gr.Column(scale=3):
+            with gr.Column(scale=12):
+                user_input = gr.Textbox(show_label=False, placeholder="Input...", lines=10, container=False)
+            with gr.Column(min_width=32, scale=1):
+                submitBtn = gr.Button("Submit")
+        with gr.Column(scale=1):
+            prompt_input = gr.Textbox(show_label=False, placeholder="Prompt", lines=10, container=False)
+            pBtn = gr.Button("Set Prompt")
+        with gr.Column(scale=1):
+            emptyBtn = gr.Button("Clear History")
+            max_length = gr.Slider(0, 10000000, value=10000, step=1.0, label="Maximum length", interactive=True)
+            top_p = gr.Slider(0, 1, value=0.8, step=0.01, label="Top P", interactive=True)
+            temperature = gr.Slider(0.01, 1, value=0.6, step=0.01, label="Temperature", interactive=True)
+
+
+    def user(query, history):
+        return "", history + [[query, ""]]
+
+
+    def set_prompt(prompt_text):
+        return [[prompt_text, "Set prompt successfully"]]
+
+
+    pBtn.click(set_prompt, inputs=[prompt_input], outputs=chatbot)
+
+    submitBtn.click(user, [user_input, chatbot], [user_input, chatbot], queue=False).then(
+        predict, [chatbot, prompt_input, max_length, top_p, temperature], chatbot
+    )
+    emptyBtn.click(lambda: (None, None), None, [chatbot, prompt_input], queue=False)
+
+demo.queue()
+demo.launch()

requirements.txt

@@ -1 +1,11 @@
-huggingface_hub==0.22.2
+gradio==4.41.0
+torch==2.2.0
+transformers==4.44.0
+spaces==0.29.2
+accelerate==0.33.0
+sentencepiece==0.2.0
+huggingface-hub==0.24.5
+sentencepiece==0.2.0
+jinja2==3.1.4
+sentence_transforme==3.0.1
+tiktoke==0.7.0