test-space / app.py
reynardryanda's picture
another test
04e4555
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()