import os
import sys
import subprocess

if not os.path.exists("ChatGLM-6b-onnx-u8s8"):
    subprocess.run(["git", "lfs", "install"])
    subprocess.run(["git", "clone", "https://huggingface.co/K024/ChatGLM-6b-onnx-u8s8"])
    os.chdir("ChatGLM-6b-onnx-u8s8")
    subprocess.run(["pip", "install", "-r", "requirements.txt"])
    sys.path.append(os.getcwd())
else:
    sys.path.append(os.path.join(os.getcwd(), "ChatGLM-6b-onnx-u8s8"))

from model import ChatGLMModel#, chat_template

model = ChatGLMModel()
# history = []

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

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

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

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

文本：
"""

def sum_chain_l1(text, p_bar):
    docs = []
    for i in p_bar(range(len(text)//2000+1)):
        t = text[i*2000:i*2000+2048]
        if len(t) > 0:
            for answer in model.generate_iterate(prompt+t,
                            max_generated_tokens=max_tokens,
                            top_k=top_k,
                            top_p=top_p,
                            temperature=temperature):
                yield answer
            docs.append(answer)
    return docs

def sum_chain_l2_deprecated(docs):
  hist = ''
  for doc in tqdm(docs):
    hist = model.response(prompt+"\n"+hist+"\n"+doc)
  return hist

import gradio as gr

def greet(x, progress=gr.Progress()):
    progress(0, desc="Starting...")
    docs = []
    for doc in sum_chain_l1(x, progress.tqdm):
        yield doc
    return '===== summarized parts ====='.join(doc)

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