import gradio as gr
import random
import time
from transformers import AutoModelForQuestionAnswering, pipeline

# 设置使用GPU
device = "cuda" if torch.cuda.is_available() else "cpu"

# 初始化模型
model_name = "bert-base-uncased"
model = AutoModelForQuestionAnswering.from_pretrained(model_name).to(device)

# 创建管道
qa_pipeline = pipeline("question-answering", model=model, tokenizer=model_name)

# 假设这是队标的URL或本地路径
team_logo_url = './team-logo.jpg'  # 替换为你的队标图片路径

speakers = ["subway"]
project_intro = """
## Automatic learning DEMO
  
**atomatic learning是一个训练方式的探索**    
**通过采样视频素材来微调LLM和音色模型**    
尽可能模仿出训练对象的说话语气和音色   
     
_训练结果来自嘉然4个视频,时长6小时_

项目使用方案     
- **语音**：bert-vits2 - **LLM**：Chat GLM-2    
特别感谢：trochkera开源训练库    


**注意**
- 所有文本数据来自视频ASR提取，价值观与基本信息未经人工对齐。所以我们不能对输出结果进行保证    
_（例如她不知道自己是女生还是男生，因为视频未提及。也可能出现LLM的“幻觉”现象）_
- 所用的BERT-VITS2版本缘故，项目暂时只支持中文    

项目即将开源，测试完成后率先打包至autodl社区-Code with GPU

_我们诚挚地感谢您的支持与反馈。_   

_feedback email:hhyxnh@gmail_


"""

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column(scale=4):
            gr.Image(value=team_logo_url, width=400)  
            gr.Markdown(project_intro)  
            
        with gr.Column(scale=8):
            chatbot = gr.Chatbot()
            msg = gr.Textbox(placeholder="Type your message here...")
            # people_speak = gr.Textbox(value="Human", label="Your name")
            bot_speak = gr.Dropdown(
                choices=speakers, value=speakers[0], label="Speaker"
            )
            with gr.Row():
                submit = gr.Button("Chat!", variant="primary")
                clear = gr.ClearButton([msg, chatbot])
    def respond(people_speak, bot_speak, message, chat_history):
        bot_message = random.choice(["How are you?", "I love you", "I'm very hungry"])
        chat_history.append(( message, bot_speak + ": " + bot_message))
        time.sleep(1)  
        return "", chat_history  

    submit.click(respond, inputs=[bot_speak, msg, chatbot], outputs=[msg, chatbot])

if __name__ == "__main__":
    demo.launch(share=True)