import os
import sys
import subprocess

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())

from model import ChatGLMModel#, chat_template

model = ChatGLMModel()
# history = []

max_tokens = 512
temperature = 1.0
top_p = 0.7
top_k = 50

from typing import Any, List, Mapping, Optional
from langchain.callbacks.manager import CallbackManagerForLLMRun
from langchain.llms.base import LLM

class CustomLLM(LLM):
    model: ChatGLMModel
    # history: List

    @property
    def _llm_type(self) -> str:
        return "custom"

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
    ) -> str:
        if stop is not None:
            raise ValueError("stop kwargs are not permitted.")
        # prompt = chat_template(self.history, prompt)
        for answer in self.model.generate_iterate(prompt,
                            max_generated_tokens=max_tokens,
                            top_k=top_k,
                            top_p=top_p,
                            temperature=temperature):
          pass

        # self.history = self.history + [(question, answer)]
        return answer

    @property
    def _identifying_params(self) -> Mapping[str, Any]:
        """Get the identifying parameters."""
        return {"model": "ChatGLMModel"}


llm = CustomLLM(model=model)

import gradio as gr
from langchain.docstore.document import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.chains.question_answering import load_qa_chain
# from langchain.embeddings import HuggingFaceEmbeddings
# from langchain.vectorstores import Chroma

# embeddings = HuggingFaceEmbeddings()
query = "總結並以點列形式舉出重點"
chain = load_qa_chain(llm, chain_type="map_reduce")
# chain = load_qa_chain(llm, chain_type="refine")

def greet(text):
    docs = [Document(page_content=text)]

    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1024, # 分割最大尺寸
        chunk_overlap=64, # 重复字数
        length_function=len
    )
    texts = text_splitter.split_documents(docs)
    # docsearch = Chroma.from_texts(texts, embeddings).as_retriever()
    # docs = docsearch.get_relevant_documents(query)
    return chain.run(input_documents=texts, question=query)

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