import os
import subprocess

if not os.path.exists("chatglm4-ggml-int4.bin"):
    os.system("wget https://huggingface.co./npc0/chatglm-4-9b-int4/resolve/main/chatglm4-ggml-int4.bin")

import chatglm_cpp
pipeline = chatglm_cpp.Pipeline("chatglm4-ggml-int4.bin")
# pipeline.chat([chatglm_cpp.ChatMessage(role="user", content="你好")])
# history = []

prompt = """
現在有些文本，文本詳細且複雜。 它包含細節,可以縮減和綜合為關鍵要點。 你的任務是提取最重要的概念,重點關注主要思路,提供一個概述而不失去精髓。 你的總結應該:

• 簡潔但足夠充分,可以代表所有重要信息
• 使用正確的句式和連貫的流程
• 捕捉誰、什麼、何時、在哪裡、為什麼和如何
• 盡可能地保留原始風格和風格
• 你必須遵循“摘要”格式:

摘要:
用2至3個句子簡要陳述主要主題和主要發現。

主要要點:
• 要點1 - 最重要的發現或細節
• 要點2 - 第二重要的觀黵
• 要點3 - 第三個重要的信息
• 要點4(可選) - 另一個要點
• 要點5(可選) - 最後的關鍵總結要點

文本：
"""

def sum_chain_l1(text, p_bar):
    docs = []
    for i in p_bar(range(len(text)//8000+1)):
        t = text[i*8000:i*8000+8196]
        if len(t) > 0:
            a = ''
            for answer in pipeline.chat(
                            [chatglm_cpp.ChatMessage(
                                role="user",
                                content=prompt+t)],
                            stream=True):
                a += answer.content
                yield f"{'='*8} {i+1}/{len(text)//8000+1} {'='*8}\n{a}"
            docs.append(answer)
    yield docs

def sum_chain_l2_deprecated(docs, p_bar):
    hist = docs[0]
    i = 0
    for doc in p_bar(docs[1:]):
        i += 1
        a = ''
        for answer in pipeline.chat(
                            [chatglm_cpp.ChatMessage(
                                role="user",
                                content=prompt+"\n"+hist+"\n"+doc)],
                            stream=True):
            a += answer.content
            yield f"{'='*8} {i}/{len(docs)} {'='*8}\n{a}"
        hist = answer
    yield hist

import gradio as gr

def greet(text, progress=gr.Progress()):
    progress(0, desc="Reading")
    docs = []
    for doc in sum_chain_l1(text, progress.tqdm):
        if isinstance(doc, str):
            yield '# drafting summary', doc
    # yield 'stage 1 finished', '\n===== summarized parts =====\n'.join(doc)
    progress(0, desc="Refinement")
    for ans in sum_chain_l2_deprecated(doc, progress.tqdm):
        if isinstance(ans, str):
            yield '# refining summary', ans
    return '# final result', ans

gr.Markdown("# Text Summarization\nStage 1 = draft -> Stage 2 = refine -> final result")

iface = gr.Interface(fn=greet,
                     inputs=gr.Textbox(lines=20,
                                       placeholder="Text Here..."),
                     outputs=["text", "text"])
iface.launch()