Spaces:
Running
Running
| import base64 | |
| from copy import deepcopy | |
| import gradio as gr | |
| import time | |
| from pathlib import Path | |
| from yuewen_api import call_audiochat, call_asr, call_tts | |
| CACHE_DIR = "/tmp/gradio/" | |
| CACHE_CLEAN_AGE = 86400 | |
| CHINESE_PROMPT_CONTENT = """你是一个为对话而设计的人工智能模型，目前无法连接到互联网。 | |
| 当你需要唱歌时，请以（哼唱）开头。当你需要rap或说唱时，请以（RAP）开头。当你需要快速说话时，请以（快速）开头。当你需要慢速说话时，请以（慢速）开头。 | |
| 现在，你需要倾听用户的语音内容，并以礼貌、简洁、口语化的文本进行回复。你需要尽量用户的语种进行回复。""" | |
| ENGLISH_PROMPT_CONTENT = """You are an AI designed for conversation, currently unable to connect to the internet. | |
| when you need to sing or rap, start your response with (RAP). when you need to speak fast, you start your response with (fast). when you need to speak fast, you start your response with (slow) | |
| Now, you need to listen to the user's voice content and respond with politely, concise, conversational text. Respond in accordance with the user's language.""" | |
| PROMPT_TEMPLATE = { | |
| "English Prompt": ENGLISH_PROMPT_CONTENT, | |
| "Chinese Prompt": CHINESE_PROMPT_CONTENT, | |
| } | |
| def make_system_prompt(prompt): | |
| return {"role": "system", "content": prompt} | |
| BASE_DIR = Path(__file__).resolve().parent | |
| CHAT_EXAMPLES = [ | |
| ["Input audio only", None, f"{BASE_DIR}/assets/tell_me_a_short_story_EN.wav"], | |
| [ | |
| "Input text only", | |
| "What did the speaker mean when they said, it's raining cats and dogs?", | |
| None, | |
| ], | |
| [ | |
| "Text and audio mixed input", | |
| "Answer the following query in English", | |
| f"{BASE_DIR}/assets/request_rap_ZH.wav", | |
| ], | |
| ] | |
| def file_to_base64(filename): | |
| with open(filename, "rb") as f: | |
| file_content = f.read() | |
| return base64.b64encode(file_content).decode("utf-8") | |
| def get_audio_format(audio_filename): | |
| extension = Path(audio_filename).suffix.lower() | |
| if extension == ".mp3": | |
| return "mp3" | |
| elif extension == ".wav": | |
| return "wav" | |
| else: | |
| return "other" | |
| def add_message(chatbot, history, mic, text): | |
| if not mic and not text: | |
| return chatbot, history, "Input is empty" | |
| content = [] | |
| if text: | |
| chatbot.append({"role": "user", "content": text}) | |
| content.append({"type": "text", "text": text}) | |
| print(f"add message {text=}") | |
| if mic and Path(mic).exists(): | |
| chatbot.append({"role": "user", "content": {"path": mic}}) | |
| content.append( | |
| { | |
| "type": "input_audio", | |
| "input_audio": { | |
| "data": file_to_base64(mic), | |
| "format": get_audio_format(mic), | |
| }, | |
| } | |
| ) | |
| print(f"add message {mic=}") | |
| history.append({"role": "user", "content": content}) | |
| return chatbot, history, None | |
| def get_tmp_audio_path(): | |
| import tempfile | |
| temp_audio = tempfile.NamedTemporaryFile(dir=CACHE_DIR, delete=False, suffix=".mp3") | |
| return temp_audio.name | |
| def predict(chatbot, history, user_prompt, enable_asr): | |
| """Generate a response from the model.""" | |
| start_time = time.time() | |
| try: | |
| messages = [make_system_prompt(user_prompt)] + history | |
| if enable_asr: | |
| asr_text = None | |
| if not isinstance(chatbot[-1]["content"], str): | |
| if isinstance(chatbot[-1]["content"], dict): | |
| user_audio_path = chatbot[-1]["content"]["path"] | |
| else: | |
| user_audio_path = chatbot[-1]["content"][0] | |
| print(f"do asr {user_audio_path=}") | |
| asr_text = call_asr(user_audio_path) | |
| print(f"{asr_text=}") | |
| if asr_text: | |
| last_input = deepcopy(history[-1]) | |
| last_input["content"].append({"type": "text", "text": asr_text}) | |
| messages = ( | |
| [make_system_prompt(user_prompt)] + history[:-1] + [last_input] | |
| ) | |
| text = call_audiochat(messages) | |
| print(f"predict {text=}") | |
| audio_path = get_tmp_audio_path() | |
| call_tts(text, audio_path) | |
| print(f"save_tmp_audio {audio_path=}") | |
| chatbot.append({"role": "assistant", "content": text}) | |
| chatbot.append({"role": "assistant", "content": {"path": audio_path}}) | |
| history.append({"role": "assistant", "content": text}) | |
| except Exception as e: | |
| import traceback | |
| traceback.print_exc() | |
| print(f"get a excption: {e=}") | |
| gr.Warning(f"get a excption: {str(e)}, try regenerate") | |
| print(f"predict cost {(time.time() - start_time):.2f}s") | |
| return chatbot, history | |
| def _launch_demo(args): | |
| with gr.Blocks(delete_cache=(3600, CACHE_CLEAN_AGE)) as demo: | |
| # 保存 chat 历史，不需要每次再重新拼格式 | |
| history = gr.State([]) | |
| gr.Markdown("""<center><font size=8>Step Audio Chat</center>""") | |
| gr.Markdown( | |
| """<font size=4>This preview demonstrates core functionalities. To unlock the cormplete real-time voice conversation system with end-to-end encryption and advanced features, download the [Yuewen APP](https://m.yuewen.cn/call-app) with the link or via QR Code.</font>""" | |
| ) | |
| with gr.Accordion(label="Click to view the QR code ", open=False): | |
| gr.Image( | |
| value="assets/yuewen.jpeg", | |
| interactive=False, | |
| show_label=False, | |
| show_download_button=False, | |
| show_fullscreen_button=False, | |
| ) | |
| with gr.Accordion( | |
| label="The performance of English prompts is not as stable as that of Chinese prompts. You can click here to change sys prompt.", | |
| open=False, | |
| ): | |
| prompt_choice = gr.Radio( | |
| choices=list(PROMPT_TEMPLATE.keys()), | |
| value=list(PROMPT_TEMPLATE.keys())[0], | |
| label="Select prompt template", | |
| ) | |
| user_prompt = gr.Textbox( | |
| show_label=False, | |
| value=list(PROMPT_TEMPLATE.values())[0], | |
| lines=6, | |
| ) | |
| prompt_choice.change( | |
| fn=lambda choice: PROMPT_TEMPLATE[choice], | |
| inputs=prompt_choice, | |
| outputs=user_prompt, | |
| ) | |
| chatbot = gr.Chatbot( | |
| elem_id="chatbot", | |
| avatar_images=["assets/user.png", "assets/assistant.png"], | |
| min_height=800, | |
| type="messages", | |
| show_share_button=True, | |
| ) | |
| mic = gr.Audio( | |
| label="Only MP3 and WAV formats are supported for audio uploads.", | |
| sources=["microphone", "upload"], | |
| type="filepath", | |
| ) | |
| text = gr.Textbox( | |
| placeholder="Enter message ...", | |
| label="Input text message", | |
| show_label=False, | |
| ) | |
| with gr.Row(): | |
| enable_asr = gr.Checkbox( | |
| value=True, | |
| label="Enhance understanding capability by ASR", | |
| ) | |
| with gr.Row(): | |
| clear_btn = gr.Button("🧹 Clear History (清除历史)") | |
| regen_btn = gr.Button("🤔️ Regenerate (重试)") | |
| gen_btn = gr.Button("🚀 Generate Response") | |
| example_comment = gr.Textbox(label="Input data type", visible=False) | |
| example_text = gr.Textbox(label="Input text message", visible=False) | |
| example_audio = gr.Audio( | |
| label="Input Audio", | |
| type="filepath", | |
| visible=False, | |
| ) | |
| def update_examples(_, text, mic, user_prompt, enable_asr): | |
| chatbot = [] | |
| history = [] | |
| chatbot, history, error = add_message(chatbot, history, mic, text) | |
| if error: | |
| gr.Warning(error) | |
| print(f"update_examples error") | |
| return chatbot, history | |
| else: | |
| chatbot, history = predict(chatbot, history, user_prompt, enable_asr) | |
| print(f"update_examples done") | |
| return chatbot, history | |
| with gr.Row(): | |
| gr.Examples( | |
| fn=update_examples, | |
| examples=CHAT_EXAMPLES, | |
| inputs=[ | |
| example_comment, | |
| example_text, | |
| example_audio, | |
| user_prompt, | |
| enable_asr, | |
| ], | |
| outputs=[chatbot, history], | |
| run_on_click=True, | |
| ) | |
| def on_submit(chatbot, history, mic, text, user_prompt, enable_asr): | |
| chatbot, history, error = add_message(chatbot, history, mic, text) | |
| if error: | |
| gr.Warning(error) | |
| return chatbot, history, None, None | |
| else: | |
| chatbot, history = predict(chatbot, history, user_prompt, enable_asr) | |
| return chatbot, history, None, None | |
| gen_btn.click( | |
| fn=on_submit, | |
| inputs=[chatbot, history, mic, text, user_prompt, enable_asr], | |
| outputs=[chatbot, history, mic, text], | |
| show_progress=True, | |
| ) | |
| def clear_chat_history(): | |
| return [], [] | |
| clear_btn.click( | |
| clear_chat_history, | |
| outputs=[chatbot, history], | |
| show_progress=True, | |
| ) | |
| def regenerate(chatbot, history, user_prompt, enable_asr): | |
| while chatbot and chatbot[-1]["role"] == "assistant": | |
| chatbot.pop() | |
| while history and history[-1]["role"] == "assistant": | |
| print(f"discard {history[-1]}") | |
| history.pop() | |
| return predict(chatbot, history, user_prompt, enable_asr) | |
| regen_btn.click( | |
| regenerate, | |
| [chatbot, history, user_prompt, enable_asr], | |
| [chatbot, history], | |
| show_progress=True, | |
| ) | |
| demo.queue().launch( | |
| share=False, | |
| server_port=args.server_port, | |
| server_name=args.server_name, | |
| ) | |
| if __name__ == "__main__": | |
| from argparse import ArgumentParser | |
| import os | |
| parser = ArgumentParser() | |
| parser.add_argument("--model-path", type=str, help="Tokenizer and TTS model path.") | |
| parser.add_argument( | |
| "--server-port", type=int, default=7860, help="Demo server port." | |
| ) | |
| parser.add_argument( | |
| "--server-name", type=str, default="0.0.0.0", help="Demo server name." | |
| ) | |
| args = parser.parse_args() | |
| _launch_demo(args) | |