Update README.md

README.md

@@ -5,7 +5,7 @@ language:
 thumbnail: "https://upload.wikimedia.org/wikipedia/commons/a/a2/Farsi.svg"
 tags:
 - entailment
-- parsbert 
+- mt5 
 - persian 
 - farsi 
 license: "CC BY-NC-SA 4.0"
@@ -22,36 +22,33 @@ This is a model for textual entailment problems.
 Here is an example of how you can run this model: 
 
 ```python 
-import torch
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
-import numpy as np
+from transformers import MT5ForConditionalGeneration, MT5Tokenizer
 
-labels = ["entails", "contradicts", "neutral"]
-model_name_or_path = "persiannlp/mbert-base-parsinlu-snli-entailment"
-model = AutoModelForSequenceClassification.from_pretrained(model_name_or_path)
-tokenizer = AutoTokenizer.from_pretrained(model_name_or_path,)
+model_size="base"
+model_name = f"persiannlp/mt5-{model_size}-parsinlu-snli-entailment"
+tokenizer = MT5Tokenizer.from_pretrained(model_name)
+model = MT5ForConditionalGeneration.from_pretrained(model_name)
 
 
-def model_predict(text_a, text_b):
-    features = tokenizer( [(text_a, text_b)], padding="max_length", truncation=True, return_tensors='pt')
-    output = model(**features)
-    logits = output[0]
-    probs = torch.nn.functional.softmax(logits, dim=1).tolist()
-    idx = np.argmax(np.array(probs))
-    print(labels[idx], probs)
+def run_model(premise, hypothesis, **generator_args):
+    input_ids = tokenizer.encode(f"{premise}<sep>{hypothesis}", return_tensors="pt")
+    res = model.generate(input_ids, **generator_args)
+    output = tokenizer.batch_decode(res, skip_special_tokens=True)
+    print(output)
+    return output
 
 
-model_predict(
+run_model(
     "این مسابقات بین آوریل و دسامبر در هیپودروم ولیفندی در نزدیکی باکرکی ، ۱۵ کیلومتری (۹ مایل) غرب استانبول برگزار می شود.",
     "در ولیفندی هیپودروم، مسابقاتی از آوریل تا دسامبر وجود دارد."
 )
 
-model_predict(
+run_model(
 "آیا کودکانی وجود دارند که نیاز به سرگرمی دارند؟",
     "هیچ کودکی هرگز نمی خواهد سرگرم شود.",
 )
 
-model_predict(
+run_model(
     "ما به سفرهایی رفته ایم که در نهرهایی شنا کرده ایم",
     "علاوه بر استحمام در نهرها ، ما به اسپا ها و سونا ها نیز رفته ایم."
 )