import os import subprocess import gradio as gr from epub2txt import epub2txt from websocket import create_connection if not os.path.exists(os.getenv("checkpoint_path")): os.system("git clone --recurse-submodules https://github.com/ztxz16/fastllm.git") os.system("cd fastllm; mkdir build; cd build; cmake ..; make -j; cd tools; python setup.py install --user --prefix=") os.system("wget https://huggingface.co/huangyuyang/chatglm2-6b-int4.flm/resolve/main/chatglm2-6b-int4.flm") subprocess.Popen("uvicorn", "api:app") class GUI: def __init__(self, *args, **kwargs): with gr.Blocks() as demo: with gr.Row(): gr.Markdown(scale=2).attach_load_event(self.hello, None) gr.LoginButton() gr.LogoutButton() out = gr.Markdown() inp = gr.File(file_types=['.epub']) inp.change(self.process, inp, out) self.ws = None self.out = [] demo.queue(concurrency_count=2).launch() def process(self, file, profile: gr.OAuthProfile | None): if profile is None: return gr.update(value='Login to access the tool.') h = '' chapter_titles = epub2txt.content_titles title = epub2txt.title if self.ws is None: self.ws = create_connection(f"ws://localhost:8000/ws") self.ws.send(file.name) res = '' while 'output: ' not in res: res = self.ws.recv() if 'chsum: ' in res: self.out.append(res.replace("chsum: ", "")) elif 'draft_sum: ' in res: h = res[11:] elif 'output: ' in res: self.ws.close() self.ws = None self.out = [] yield gr.update(value=res) yield gr.update( value=f"# {title}\n\n" +h+ "\n\n".join( [f"## {ct}\n\n{c}" for ct, c in zip(chapter_titles, self.out)])) def hello(self, profile: gr.OAuthProfile | None): if profile is None: return '# ePub summarization tool\n\nLogin to access the tool.' return f"# ePub summarization tool\n\nWelcome {profile.name}!!" GUI()