import gradio as gr
import re
from transformers import AutoTokenizer, AutoModel

MODEL_NAME = "THUDM/chatglm-6b-int4"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
model = AutoModel.from_pretrained(MODEL_NAME, trust_remote_code=True).half().cpu()

def summarize(transcript, sentence_count):
    history = []
    
    prompt = f"""视频脚本：{transcript}。我希望你作为一名专业的视频内容编辑，帮我用中文总结视频脚本的内容精华。请先用一句简短的话总结视频梗概。然后再请你将视频字幕文本进行总结（字幕中可能有错别字，如果你发现了错别字请改正）。请你以无序列表的方式返回，请注意不要超过{sentence_count}条哦，确保所有的句子都足够精简，清晰完整，祝你好运！"""

    response, history = model.chat(tokenizer, prompt, history=history)

    return response


demo = gr.Interface(fn = summarize,
                    inputs = [gr.Textbox(lines=10, 
                                                placeholder="Input something...",
                                                label='Text here !!'),
                                gr.Slider(minimum=1,
                                                    maximum=10,
                                                    step=1,
                                                    label='Sentence Count')],
                    outputs = [gr.Textbox(lines=10,
                                                    label="Summary")],
                        
                    title = "🎈 Summarizer 🎈")

demo.launch()