Create app.py

app.py

@@ -0,0 +1,72 @@
+import gradio as gr
+from llamafactory.chat import ChatModel
+from llamafactory.extras.misc import torch_gc
+import re
+import space
+
+def split_into_sentences(text):
+    sentence_endings = re.compile(r'(?<!\w\.\w.)(?<![A-Z][a-z]\.)(?<=\.|\?|\!)\s')
+    sentences = sentence_endings.split(text)
+    return [sentence.strip() for sentence in sentences if sentence]
+
+@space.GPU
+def process_paragraph(paragraph, progress=gr.Progress()):
+    sentences = split_into_sentences(paragraph)
+    results = []
+    total_sentences = len(sentences)
+    for i, sentence in enumerate(sentences):
+        progress((i + 1) / total_sentences)
+        messages.append({"role": "user", "content": sentence})
+        sentence_response = ""
+        for new_text in chat_model.stream_chat(messages, temperature=0.7, top_p=0.9, top_k=50, max_new_tokens=300):
+            sentence_response += new_text.strip()
+        category = sentence_response.strip().lower().replace(' ', '_')
+        if category != "fair":
+            results.append((sentence, category))
+        else:
+            results.append((sentence, "fair"))
+        messages.append({"role": "assistant", "content": sentence_response})
+        torch_gc()
+    return results
+
+args = dict(
+  model_name_or_path="princeton-nlp/Llama-3-Instruct-8B-SimPO",  # 使用量化的 Llama-3-8B-Instruct 模型
+  adapter_name_or_path="StevenChen16/llama3-8b-compliance-review",                 # 加载保存的 LoRA 适配器
+  template="llama3",                                      # 与训练时使用的模板相同
+  finetuning_type="lora",                                 # 与训练时使用的微调类型相同
+  quantization_bit=8,                                     # 加载 4-bit 量化模型
+  use_unsloth=True,                                       # 使用 UnslothAI 的 LoRA 优化以加速生成
+)
+chat_model = ChatModel(args)
+messages = []
+
+# 定义类型到颜色的映射
+label_to_color = {
+    "fair": "green",
+    "limitation_of_liability": "red",
+    "unilateral_termination": "orange",
+    "unilateral_change": "yellow",
+    "content_removal": "purple",
+    "contract_by_using": "blue",
+    "choice_of_law": "cyan",
+    "jurisdiction": "magenta",
+    "arbitration": "brown",
+}
+
+with gr.Blocks() as demo:
+    
+    with gr.Row(equal_height=True):
+        with gr.Column():
+            input_text = gr.Textbox(label="Input Paragraph", lines=10, placeholder="Enter the paragraph here...")
+            btn = gr.Button("Process")
+        with gr.Column():
+            output = gr.HighlightedText(label="Processed Paragraph", color_map=label_to_color)
+            progress = gr.Progress()
+
+    def on_click(paragraph):
+        results = process_paragraph(paragraph, progress=progress)
+        return results
+
+    btn.click(on_click, inputs=input_text, outputs=[output])
+
+demo.launch()