import gradio as gr
import torch
import numpy as np
from transformers import BertTokenizer, BertModel

# 1. 核心引擎加载：强制使用慢速分词模式以兼容专利模型
model_name = "anferico/bert-for-patents"

print("正在加载专利专用分词器...")
tokenizer = BertTokenizer.from_pretrained(
    model_name, 
    use_fast=False, 
    local_files_only=False
)

print("正在加载 BERT 专利模型主体 (0.3B)...")
model = BertModel.from_pretrained(model_name)
model.eval() # 设置为评估模式

def get_embedding(text):
    if not text.strip():
        return [0] * 1024
    try:
        # 将文本转化为 1024 维的专利语义向量
        inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
        with torch.no_grad():
            outputs = model(**inputs)
        # 取最后一层隐藏状态的平均值作为特征向量 (Mean Pooling)
        embeddings = outputs.last_hidden_state.mean(dim=1).squeeze().numpy()
        return embeddings.tolist()
    except Exception as e:
        return f"向量化错误: {str(e)}"

def compute_similarity(text1, text2):
    if not text1.strip() or not text2.strip():
        return "请输入完整的对比文本"
    try:
        v1 = np.array(get_embedding(text1))
        v2 = np.array(get_embedding(text2))
        # 计算余弦相似度
        similarity = np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))
        return f"{similarity:.4f}"
    except Exception as e:
        return f"计算失败: {str(e)}"

# 2. OpenIPOS 旗舰版界面装修
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.HTML("""
        <div style="display: flex; align-items: center; gap: 20px; padding: 10px;">
            <img src="https://s3.bmp.ovh/2026/03/30/rIAGKg0O.png" style="height: 50px; width: auto;">
            <div>
                <h1 style="margin: 0; font-size: 24px;">OpenIPOS Patent Vector Engine</h1>
                <p style="margin: 0; color: #666;">基于谷歌 BERT-for-Patents 的专业语义标尺 (0.3B)</p>
            </div>
        </div>
    """)
    
    with gr.Tabs():
        with gr.TabItem("🔍 语义相似度比对"):
            with gr.Row():
                with gr.Column():
                    t1 = gr.Textbox(label="本申请技术特征 (权利要求)", lines=8, placeholder="请输入权利要求书中的核心技术特征...")
                with gr.Column():
                    t2 = gr.Textbox(label="对比文件技术特征 (D1/D2)", lines=8, placeholder="请输入对比文件中的对应描述内容...")
            
            compare_btn = gr.Button("开始语义比对", variant="primary")
            similarity_output = gr.Label(label="语义重合度评分 (0-1)")
            
            gr.Markdown("""
            ### ⚖️ 判读参考：
            * **0.85 - 1.00**：极高风险。语义高度重合，可能构成实质性相同。
            * **0.70 - 0.85**：高风险。核心特征重叠，建议重点修改权利要求。
            * **0.50 - 0.70**：中等风险。存在语义关联，具有一定的创造性抗辩空间。
            """)
            
            compare_btn.click(compute_similarity, inputs=[t1, t2], outputs=similarity_output)

        with gr.TabItem("📊 向量化测试 (API)"):
            input_text = gr.Textbox(label="输入专利描述", lines=5)
            v_btn = gr.Button("生成 1024 维特征向量")
            vector_out = gr.JSON(label="向量数组")
            v_btn.click(get_embedding, inputs=input_text, outputs=vector_out)

    gr.Markdown("---")
    gr.Markdown("© 2026 OpenIPOS Global Limited. 核心算法：BERT for Patents | 算力驱动：Hugging Face CPU Basic")

# 启动
demo.launch()