Update README.md

README.md

@@ -23,17 +23,90 @@ It achieves the following results on the evaluation set:
 
 ## Model description
 
-More information needed
-
-## Intended uses & limitations
-
-More information needed
-
-## Training and evaluation data
-
-More information needed
-
-## Training procedure
+Machine Translation model from Hindi to English on bart small model.
+
+## Inference and evaluation
+
+```python
+import torch
+import evaluate
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+class BartSmall():
+    def __init__(self, model_path = 'ar5entum/bart_hin_eng_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)
+
+    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)
+
+        return all_predictions
+
+model = BartSmall(device='cuda')
+
+input_texts = [
+    "यह शोध्य रकम है।", 
+    "जानने के लिए देखें ये वीडियो.",
+    "वह दो बेटियों व एक बेटे का पिता था।"
+    ]
+ground_truths = [
+    "This is a repayable amount.",
+    "Watch this video to find out.",
+    "He was a father of two daughters and a son."
+    ]
+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:  1.2374696731567383
+# ‾‾‾‾‾‾‾‾‾‾‾‾
+# Input text:	 यह शोध्य रकम है।
+# Prediction:	 This is a repayable amount.
+# Ground Truth:	 This is a repayable amount.
+# ‾‾‾‾‾‾‾‾‾‾‾‾
+# Input text:	 जानने के लिए देखें ये वीडियो.
+# Prediction:	 View these videos to know.
+# Ground Truth:	 Watch this video to find out.
+# ‾‾‾‾‾‾‾‾‾‾‾‾
+# Input text:	 वह दो बेटियों व एक बेटे का पिता था।
+# Prediction:	 He was a father of two daughters and a son.
+# Ground Truth:	 He was a father of two daughters and a son.
+# {'bleu': 0.747875245486914, 'precisions': [0.8260869565217391, 0.75, 0.7647058823529411, 0.7857142857142857], 'brevity_penalty': 0.9574533680683809, 'length_ratio': 0.9583333333333334, 'translation_length': 23, 'reference_length': 24}
+```
 
 ### Training hyperparameters