import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

# 確実に存在する軽量日本語モデル
model_name = "rinna/japanese-gpt2-small"
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
model = AutoModelForCausalLM.from_pretrained(model_name)

# 採点用プロンプト
PROMPT_TEMPLATE = """
あなたは作文採点の教師です。Q{question_number}の回答を0〜12点で採点してください。
回答: {answer}
採点ルール: 論理性、文字数、日本語の正確さ、他の比較や妄想の程度などを考慮し0〜12点で評価。
"""

# 採点関数
def grade(question_number, answer):
    if not question_number:
        return "Q77またはQ78の回答ですか？"
    
    prompt = PROMPT_TEMPLATE.format(question_number=question_number, answer=answer)
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs, max_new_tokens=20, temperature=0.5)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# 作文採点チャットボット (0〜12点)")
    qnum = gr.Radio(["77", "78"], label="問題番号")
    text = gr.Textbox(label="回答", lines=6, placeholder="ここに回答を入力してください")
    output = gr.Textbox(label="採点結果")
    btn = gr.Button("採点する")
    btn.click(grade, inputs=[qnum, text], outputs=output)

demo.launch()