taprosoft commited on
Commit
1e92b0d
·
verified ·
1 Parent(s): e9b60f2

Update backends/smoldocling.py

Browse files
Files changed (1) hide show
  1. backends/smoldocling.py +2 -2
backends/smoldocling.py CHANGED
@@ -12,7 +12,7 @@ from PIL import Image
12
  from transformers import AutoModelForVision2Seq, AutoProcessor
13
 
14
  DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
15
- MAX_PAGES = 2
16
 
17
  # Initialize processor and model
18
  processor = AutoProcessor.from_pretrained("ds4sd/SmolDocling-256M-preview")
@@ -49,7 +49,7 @@ def convert_smoldocling(path: str, file_name: str):
49
  inputs = inputs.to(DEVICE)
50
 
51
  # Generate outputs
52
- generated_ids = model.generate(**inputs, max_new_tokens=8192 * 2)
53
  prompt_length = inputs.input_ids.shape[1]
54
  trimmed_generated_ids = generated_ids[:, prompt_length:]
55
  doctags = processor.batch_decode(
 
12
  from transformers import AutoModelForVision2Seq, AutoProcessor
13
 
14
  DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
15
+ MAX_PAGES = 1
16
 
17
  # Initialize processor and model
18
  processor = AutoProcessor.from_pretrained("ds4sd/SmolDocling-256M-preview")
 
49
  inputs = inputs.to(DEVICE)
50
 
51
  # Generate outputs
52
+ generated_ids = model.generate(**inputs, max_new_tokens=8192)
53
  prompt_length = inputs.input_ids.shape[1]
54
  trimmed_generated_ids = generated_ids[:, prompt_length:]
55
  doctags = processor.batch_decode(