Update README.md

README.md

@@ -28,17 +28,90 @@ It achieves the following results on the evaluation set:
 
 ## Model description
 
-More information needed
+Machine Translation model from English to Hindi on bart small model.
 
-## Intended uses & limitations
+## Inference and Evaluation
 
-More information needed
+```python
+import torch
+import evaluate
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-## Training and evaluation data
+class BartSmall():
+    def __init__(self, model_path = 'ar5entum/bart_eng_hin_mt', device = None):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
+        if not device:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.device = device
+        self.model.to(device)
 
-More information needed
+    def predict(self, input_text):
+        inputs = self.tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True).to(self.device)
+        pred_ids = self.model.generate(inputs.input_ids, max_length=512, num_beams=4, early_stopping=True)
+        prediction = self.tokenizer.decode(pred_ids[0], skip_special_tokens=True)
+        return prediction
+    
+    def predict_batch(self, input_texts, batch_size=32):
+        all_predictions = []
+        for i in range(0, len(input_texts), batch_size):
+            batch_texts = input_texts[i:i+batch_size]
+            inputs = self.tokenizer(batch_texts, return_tensors="pt", max_length=512, 
+                                    truncation=True, padding=True).to(self.device)
+            
+            with torch.no_grad():
+                pred_ids = self.model.generate(inputs.input_ids, 
+                                               max_length=512, 
+                                               num_beams=4, 
+                                               early_stopping=True)
+            
+            predictions = self.tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
+            all_predictions.extend(predictions)
 
-## Training procedure
+        return all_predictions
+
+model = BartSmall(device='cuda')
+
+input_texts = [
+    "This is a repayable amount.",
+    "Watch this video to find out.",
+    "He was a father of two daughters and a son."
+    ]
+ground_truths = [
+    "यह शोध्य रकम है।", 
+    "जानने के लिए देखें ये वीडियो.",
+    "वह दो बेटियों व एक बेटे का पिता था।"
+    ]
+import time
+start = time.time()
+
+predictions = model.predict_batch(input_texts, batch_size=len(input_texts))
+end = time.time()
+print("TIME: ", end-start)
+for i in range(len(input_texts)):
+    print("‾‾‾‾‾‾‾‾‾‾‾‾")
+    print("Input text:\t", input_texts[i])
+    print("Prediction:\t", predictions[i])
+    print("Ground Truth:\t", ground_truths[i])
+bleu = evaluate.load("bleu")
+results = bleu.compute(predictions=predictions, references=ground_truths)
+print(results)
+
+# TIME:  3.65848970413208
+# ‾‾‾‾‾‾‾‾‾‾‾‾
+# Input text:	 This is a repayable amount.
+# Prediction:	 यह एक चुकौती राशि है। 
+# Ground Truth:	 यह शोध्य रकम है।
+# ‾‾‾‾‾‾‾‾‾‾‾‾
+# Input text:	 Watch this video to find out.
+# Prediction:	 इस वीडियो को बाहर ढूंढने के लिए इस वीडियो को देख�
+# Ground Truth:	 जानने के लिए देखें ये वीडियो.
+# ‾‾‾‾‾‾‾‾‾‾‾‾
+# Input text:	 He was a father of two daughters and a son.
+# Prediction:	 वह दो बेटियों और एक पुत्र के पिता थे। 
+# Ground Truth:	 वह दो बेटियों व एक बेटे का पिता था।
+# {'bleu': 0.0, 'precisions': [0.4, 0.13636363636363635, 0.05263157894736842, 0.0], 'brevity_penalty': 1.0, 'length_ratio': 1.25, 'translation_length': 25, 'reference_length': 20}
+```
 
 ### Training hyperparameters