import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForCausalLM
import gradio as gr

# 模型名称
model_name = "baidu/ERNIE-4.5-0.3B-PT"

# 加载 tokenizer 和模型（首次运行可能较慢）
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    trust_remote_code=True,
    torch_dtype=torch.float32,
    device_map="auto"
)
embedding_layer = model.get_input_embeddings()

# 提取句子的平均 embedding
def get_sentence_embedding(text):
    inputs = tokenizer(text, return_tensors="pt", add_special_tokens=True)
    input_ids = inputs["input_ids"]
    with torch.no_grad():
        embeddings = embedding_layer(input_ids)  # shape: [1, seq_len, hidden_size]
        sentence_embedding = embeddings.mean(dim=1)  # shape: [1, hidden_size]
    return sentence_embedding

# Gradio 回调函数
def calculate_similarity(sentence1, sentence2):
    emb1 = get_sentence_embedding(sentence1)
    emb2 = get_sentence_embedding(sentence2)
    similarity = F.cosine_similarity(emb1, emb2).item()
    return f"Similarity: {similarity:.4f}"

# Gradio 界面
title = "Calculate two sentences's similarity by ERNIE 4.5-0.3B's embedding layer"
demo = gr.Interface(
    fn=calculate_similarity,
    inputs=[
        gr.Textbox(label="Sentence 1", placeholder="我爱北京"),
        gr.Textbox(label="Sentence 2", placeholder="我爱上海")
    ],
    outputs=gr.Textbox(label="Similarity"),
    title=title,
    description="This app uses the embedding layer of Baidu ERNIE-4.5-0.3B-PT model to compute the cosine similarity between two sentences.",
)

# 启动 Gradio app
if __name__ == "__main__":
    demo.launch()