Update README.md

README.md

@@ -21,18 +21,15 @@ Use the code below to get started with the model.
 ```python
 from PIL import Image
 import torch
-from transformers import (
-    NougatProcessor,
-    VisionEncoderDecoderModel,
-)
-
-
-processor = NougatProcessor.from_pretrained("MohamedRashad/arabic-small-nougat")
-model = VisionEncoderDecoderModel.from_pretrained("MohamedRashad/arabic-small-nougat")
+from transformers import NougatProcessor, VisionEncoderDecoderModel
 
+# Load the model and processor
+processor = NougatProcessor.from_pretrained("facebook/nougat-small")
+model = VisionEncoderDecoderModel.from_pretrained("arabic-small-nougat2")
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
+context_length = 2048
 
 def predict(image):
     # prepare PDF image for the model
@@ -43,7 +40,7 @@ def predict(image):
     outputs = model.generate(
         pixel_values.to(device),
         min_length=1,
-        max_new_tokens=2048,
+        max_new_tokens=context_length,
         bad_words_ids=[[processor.tokenizer.unk_token_id]],
     )