Update README.md

README.md

@@ -83,6 +83,55 @@ Use the code below to get started with the model.
 
 ### Training Procedure
 
+
+```python
+
+from transformers import MBart50TokenizerFast, MBartForConditionalGeneration
+import torch
+
+repo_id = "frankmorales2020/kkadian-to-spanish-translator"
+tokenizer = MBart50TokenizerFast.from_pretrained(repo_id)
+model = MBartForConditionalGeneration.from_pretrained(repo_id)
+
+def test_model(sentences):
+    tokenizer.src_lang = "[akk_AK]"
+    es_id = tokenizer.convert_tokens_to_ids("es_XX")
+    
+    for text in sentences:
+        inputs = tokenizer(text, return_tensors="pt")
+        with torch.no_grad():
+            generated_tokens = model.generate(
+                **inputs,
+                forced_bos_token_id=es_id,
+                max_new_tokens=60,
+                num_beams=5
+            )
+        translation = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+        
+        print("-" * 40)
+        print(f"AKKADIAN: {text}")
+        print(f"SPANISH:  {translation}")
+
+# List your test cases here
+examples = [
+    "šarrum bītam iṣbat",
+    "ekallam īpuš"
+]
+
+test_model(examples)
+print("-" * 40)
+```
+```python
+
+----------------------------------------
+AKKADIAN: šarrum bītam iṣbat
+SPANISH:  el rey tomó la casa
+----------------------------------------
+AKKADIAN: ekallam īpuš
+SPANISH:  él construyó el palacio
+----------------------------------------
+
+```
 <!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
 
 #### Preprocessing [optional]