import gradio as gr
from sentence_transformers import SentenceTransformer, util

model_name = "cross-encoder/ms-marco-TinyBERT-L-6"
model = SentenceTransformer(model_name)

def classify_text(input_text):
    premise = "The cat is on the mat."
    input_embedding = model.encode([input_text, premise], convert_to_tensor=True)
    similarity_score = util.pytorch_cos_sim(input_embedding[0], input_embedding[1])[0][0]
    if similarity_score > 0.7:
        return "entailment"
    elif similarity_score < 0.3:
        return "contradiction"
    else:
        return "neutral"

iface = gr.Interface(fn=classify_text, inputs="text", outputs="text", title="Cross-Encoder with SentenceTransformer")
iface.launch()