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

# Load model + tokenizer with trust_remote_code
model_name = "nikhil061307/contrastive-learning-bert-added-token-v5"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModel.from_pretrained(model_name, trust_remote_code=True)

def encode(text):
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    with torch.no_grad():
        outputs = model(**inputs)
        embedding = outputs.last_hidden_state.mean(dim=1)
    return embedding

def compute_similarity(sentence1, sentence2):
    emb1 = encode(sentence1)
    emb2 = encode(sentence2)
    similarity = F.cosine_similarity(emb1, emb2).item()
    return {"similarity": round(similarity, 4)}

demo = gr.Interface(
    fn=compute_similarity,
    inputs=[
        gr.Textbox(lines=2, placeholder="Enter first sentence..."),
        gr.Textbox(lines=2, placeholder="Enter second sentence...")
    ],
    outputs="json",
    title="Contrastive Learning BERT Similarity",
    description="Enter two sentences to compute semantic similarity using your fine-tuned BERT model."
)

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