import gradio as gr
import torch
from sentence_transformers import SentenceTransformer
from torch.nn.functional import cosine_similarity

# モデルの読み込み
model = SentenceTransformer("Shuu12121/CodeCloneDetection-ModernBERT-Owl")
model.eval()

# 閾値設定（安定性の高い0.9推奨）
THRESHOLD = 0.9


def detect_clone(code1, code2):
    if not code1.strip() or not code2.strip():
        return "❌ どちらのコードも入力してください", ""

    with torch.no_grad():
        embeddings = model.encode([code1, code2], convert_to_tensor=True)
        sim_score = cosine_similarity(embeddings[0].unsqueeze(0), embeddings[1].unsqueeze(0)).item()

    result = (
        f"🟢 類似度: {sim_score:.4f}\n→ これらのコードは **クローン** と判定されます。"
        if sim_score >= THRESHOLD
        else f"🔴 類似度: {sim_score:.4f}\n→ これらのコードは **クローンではありません**。"
    )
    return result, sim_score

# Gradioインターフェースの作成
demo = gr.Interface(
    fn=detect_clone,
    inputs=[
        gr.Textbox(label="コードスニペット1", lines=10, placeholder="例: def add(a, b): return a + b"),
        gr.Textbox(label="コードスニペット2", lines=10, placeholder="例: def sum(x, y): return x + y"),
    ],
    outputs=[
        gr.Markdown(label="判定結果"),
        gr.Number(label="Cosine Similarity")
    ],
    title="Code Clone Detection with ModernBERT-Owl 🦉",
    description="Shuu12121/CodeModernBERT-Owl によって構築された Sentence-BERT モデルを使用し、コードクローンを検出します。"
)

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