import gradio as gr
from cijiang import CiJiangRhymer
from cijiang.utils import print_results

rhymer = CiJiangRhymer(strict=False, tone=True)

def get_results(input_text, target_rhyme, beam_width=20, num_candidates=4000):
    out = rhymer.get_rhymes(text=input_text, target_rhyme=target_rhyme, beam_width=beam_width, num_candidates=num_candidates)
    mask_count = text.count("[M]")
    context = text.split('[M]')[0]

    all_outputs = f"======= 韵脚: |{target_rhyme}|\n"
    for i, (seq, log_prob) in enumerate(out[:top_results]): 
        rhymes = seq[-mask_count:].split()
        colored_rhymes = [Fore.RED + part + Style.RESET_ALL if idx < mask_count else part for idx, part in enumerate(rhymes)]
        colored_rhymes = ''.join(colored_rhymes)  # Join the parts back together

        all_outputs += f"{i+1}. {context}{colored_rhymes} (score: {log_prob:.3f})\n"
    all_outputs += "=" * 40
    return all_out

# rhymer = CiJiangRhymer(strict=False, tone=True)

# # 使用韵脚的汉字押韵
# print_results(rhymer, text="[M][M][M][M]", target_rhyme="乱比高低")

# # 使用韵脚的拼音押韵
# print_results(rhymer, text="[M][M][M][M]", target_rhyme="luan4 bi3 gao1 di1")

# # 允许声调不同
# rhymer.tone = False
# print_results(rhymer, text="[M][M][M]", target_rhyme="收视率")
# rhymer.tone = True

# # 上下文押韵
# print_results(rhymer, text="我总是天亮了[M][M]", target_rhyme="排队")
# print_results(rhymer, text="我们两个人的verse价值一辆[M][M][M][M]", target_rhyme="韩国低级")

# # # 上下文非词汇押韵
# print_results(rhymer, text="报名嘻[M][M]", target_rhyme="八个")
# # 长韵脚需要更广的搜索空间
# print_results(rhymer, text="我还是觉得音乐的[M][M][M][M][M]", target_rhyme="背靠背钻研")
# print_results(rhymer, text="我还是觉得音乐的[M][M][M][M][M]", target_rhyme="背靠背钻研", beam_width=300)
demo = gr.Interface(fn=get, inputs="text", outputs="text")
demo.launch()