from vllm import LLM, SamplingParams
import gradio as gr

llm = LLM(model="EQUES/JPharmatron-7B")
sampling_params = SamplingParams(temperature=0.0, max_tokens=512)

def compare_documents(doc_a, doc_b):
    prompt = f"""
以下の2つの文書の間にある齟齬点（食い違い）を抽出してください。
内容の差異、言い回しの変更、意味が変わる部分に注目して簡潔にリスト化してください。

文書A:
{doc_a}

文書B:
{doc_b}

齟齬点：
"""

    outputs = llm.generate(prompt, sampling_params)
    res = outputs[0].outputs[0].text
    res = res.split("\n\n")[0]
    return res.strip()


# サンプル文挿入用の関数
def load_sample():
    sample_text_a = "有効成分は37℃で30分間加熱処理した後、ろ過を行うこと。"
    sample_text_b = "加熱処理後、ろ過を行う。加熱温度は35℃とする。"
    return sample_text_a, sample_text_b

with gr.Blocks() as demo:
    gr.Markdown("## 📝 文書齟齬チェックAI（Gradioデモ）")

    with gr.Row():
        sample_button = gr.Button("サンプル1を利用")

    with gr.Row():
        input_a = gr.Textbox(label="文書A（元文書）", lines=10, placeholder="ここに原文を入力")
        input_b = gr.Textbox(label="文書B（比較対象）", lines=10, placeholder="ここに修正版を入力")

    output = gr.Textbox(label="🔍 抽出された齟齬点", lines=10)

    with gr.Row():
        run_button = gr.Button("差分を確認する")

    run_button.click(fn=compare_documents, inputs=[input_a, input_b], outputs=output)
    sample_button.click(fn=load_sample, inputs=[], outputs=[input_a, input_b])

demo.launch()