Local pytorch model does not work

Models
1

Im trying to save to local this model: AdamCodd/donut-receipts-extract · Hugging Face

Seems like I barely got it running this code (using torchScript):

import torch
import re
from PIL import Image
from transformers import DonutProcessor, VisionEncoderDecoderModel


model_name = "AdamCodd/donut-receipts-extract"
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
processor = DonutProcessor.from_pretrained(model_name)
model = VisionEncoderDecoderModel.from_pretrained(model_name)
model.to(device)

image_path = "./imagen.jpg"
image = Image.open(image_path).convert("RGB")
pixel_values = processor(image, return_tensors="pt").pixel_values



pixel_values = pixel_values.to(device)

# Generate output using model
model.eval()
with torch.no_grad():
    task_prompt = "<s_receipt>" # <s_cord-v2> for v1
    decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids
    decoder_input_ids = decoder_input_ids.to(device)



model = VisionEncoderDecoderModel.from_pretrained(model_name, torchscript=True)
traced_model = torch.jit.trace(model, (pixel_values, decoder_input_ids))


torch.jit.save(traced_model, "receipts_model.pt")

Once It generates that ‘receipts_model.pt’ file… How should I prove this model?

I have tried with this script but do not know how to make that ‘token_ids’ could return more than just one:

import torch
import re
from PIL import Image
from transformers import DonutProcessor

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# Cargar el procesador
processor = DonutProcessor.from_pretrained("AdamCodd/donut-receipts-extract")

# Cargar el modelo trazado localmente
model_path = "receipts_model.pt"  # Reemplaza con la ruta a tu modelo TorchScript
model = torch.jit.load(model_path)
model.to(device)

def load_and_preprocess_image(image_path: str, processor):
    """
    Load an image and preprocess it for the model.
    """
    image = Image.open(image_path).convert("RGB")
    pixel_values = processor(image, return_tensors="pt").pixel_values
    return pixel_values

def generate_text_from_image(model, image_path: str, processor, device):
    """
    Generate text from an image using the trained model.
    """
    # Load and preprocess the image
    pixel_values = load_and_preprocess_image(image_path, processor)
    pixel_values = pixel_values.to(device)

    # Generate output using model
    model.eval()
    with torch.no_grad():
        task_prompt = "<s_receipt>"  # <s_cord-v2> for v1
        decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids
        decoder_input_ids = decoder_input_ids.to(device)
        generated_outputs = model(pixel_values, decoder_input_ids)

    # Extract logits and convert to token ids
    logits = generated_outputs[0] if isinstance(generated_outputs, tuple) else generated_outputs
    token_ids = torch.argmax(logits, dim=-1)
    
    # Decode generated output
    decoded_text = processor.batch_decode(token_ids)[0]
    decoded_text = decoded_text.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
    decoded_text = re.sub(r"<.*?>", "", decoded_text, count=1).strip()  # remove first task start token
    decoded_text = processor.token2json(decoded_text)
    return decoded_text

# Example usage
image_path = "./imagen.jpg"  # Reemplaza con la ruta a tu imagen
extracted_text = generate_text_from_image(model, image_path, processor, device)
print("Extracted Text:", extracted_text)

Hope anyone could help me! thanks!