Spaces:
Sleeping
Sleeping
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() | |