Spaces:
Sleeping
Sleeping
| 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() |