| # Nougat for Swedish | |
| Fine-tuned version of facebook/nougat-small. Primarily trained on mathematical expressions and Swedish text. | |
| Related code at: https://github.com/will-berg/nougat-swe | |
| # Inference | |
| Code for inference on a pdf is in the GitHub repo, but can be done quickly on a provided image (of a pdf page) using the following code: | |
| ```python | |
| from PIL import Image | |
| from transformers import NougatProcessor, VisionEncoderDecoderModel | |
| import torch | |
| processor = NougatProcessor.from_pretrained("powow/nougat-swe") | |
| model = VisionEncoderDecoderModel.from_pretrained("powow/nougat-swe") | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| model.to(device) | |
| image = Image.open(image_path) | |
| pixel_values = processor(image, return_tensors="pt").pixel_values | |
| outputs = model.generate( | |
| pixel_values.to(device), | |
| min_length=1, | |
| max_new_tokens=3584, | |
| bad_words_ids=[[processor.tokenizer.unk_token_id]], | |
| ) | |
| sequence = processor.batch_decode(outputs, skip_special_tokens=True)[0] | |
| sequence = processor.post_process_generation(sequence, fix_markdown=False) | |
| print(sequence) | |
| ``` | |