another test

app.py

@@ -1,26 +1,19 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import torch
+from sentence_transformers import SentenceTransformer, util
 
-model_name = "cross-encoder/multi-nli-xlm-r-100"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
+model_name = "cross-encoder/ms-marco-TinyBERT-L-6"
+model = SentenceTransformer(model_name)
 
-def generate_prediction(input_text):
-    input_ids = tokenizer.encode(input_text, truncation=True, padding=True, return_tensors='pt')
-    outputs = model(input_ids)
-    predicted_label = torch.argmax(outputs.logits)
-    label_map = {0: "entailment", 1: "neutral", 2: "contradiction"}
-    predicted_label_text = label_map[predicted_label.item()]
-    return predicted_label_text
+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"
 
-input_text = gr.inputs.Textbox(label="Input text")
-output_text = gr.outputs.Textbox(label="Output text")
-
-gr.Interface(
-    generate_prediction,
-    inputs=input_text,
-    outputs=output_text,
-    title="Text Classifier",
-    description="A Hugging Face cross-encoder model for text classification.",
-).launch()
+iface = gr.Interface(fn=classify_text, inputs="text", outputs="text", title="Cross-Encoder with SentenceTransformer")
+iface.launch()

requirements.txt

@@ -1,2 +1 @@
-transformers
-torch
+sentence_transformers