BookSumBeta / app.py
npc0's picture
Update app.py
4b96398
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()