Update README.md

README.md

@@ -33,7 +33,29 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 ```
 
-For a badly OCRized historical text:
+And afterwards inference can be run like this:
+
+```python
+# Function to generate text
+def ocr_correction(prompt, max_new_tokens=600):
+
+    prompt = f"""### Text ###\n{prompt}\n\n\n### Correction ###\n"""
+    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
+
+    # Generate text
+    output = model.generate(input_ids,
+                            max_new_tokens=max_new_tokens,
+                            pad_token_id=tokenizer.eos_token_id,
+                            top_k=50)
+
+    # Decode and return the generated text
+    return tokenizer.decode(output[0], skip_special_tokens=True)
+
+ocr_result = ocr_correction(prompt)
+print(ocr_result)
+```
+
+A badly OCRized historical text:
 
 > NHW JICHSKV liujislatpki:.
 > 
@@ -61,29 +83,7 @@ For a badly OCRized historical text:
 > lature. And in this feeling you will, I am sure,
 > fully participate.
 
-Inference could be run like this:
-
-```python
-# Function to generate text
-def ocr_correction(prompt, max_new_tokens=600):
-
-    prompt = f"""### Text ###\n{prompt}\n\n\n### Correction ###\n"""
-    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
-
-    # Generate text
-    output = model.generate(input_ids,
-                            max_new_tokens=max_new_tokens,
-                            pad_token_id=tokenizer.eos_token_id,
-                            top_k=50)
-
-    # Decode and return the generated text
-    return tokenizer.decode(output[0], skip_special_tokens=True)
-
-ocr_result = ocr_correction(prompt)
-print(ocr_result)
-```
-
-And yield this result:
+would yield this result:
 
 > The Legislature of New Jersey assembled at Trenton, pursuant to an adjournment, on Tuesday. Both houses were organized for business, of which fact they informed the Governor, when they received the following special message.
 >