Upload inference.py with huggingface_hub

inference.py

@@ -0,0 +1,269 @@
+import torch
+import torch.nn.functional as F
+import cv2
+from PIL import Image
+from torch.utils.data import DataLoader
+from .dataset import Makeset
+
+def generate(model, image, bos_id, eos_id, device="cuda", temp=0.5, max_iter=64, penalty=1.15, top_k=5):
+    """
+    Prediction function using greedy search / top-k sampling with repetition penalty.
+    
+    Args:
+        model: SOCRATE model instance.
+        image (Tensor): Single image tensor [1, C, H, W].
+        bos_id (int): Begin-of-sequence token ID.
+        eos_id (int): End-of-sequence token ID.
+        device (str): Target device. Default: "cuda".
+        temp (float): Temperature for sampling. Lower = more greedy. Default: 0.5.
+        max_iter (int): Max number of tokens to generate. Default: 64.
+        penalty (float): Repetition penalty applied to already-seen tokens. Default: 1.15.
+        top_k (int): Number of top candidates to sample from at each step. Default: 5.
+    """
+    model.eval()
+    current_text = [bos_id]
+    generated = []
+    already_seen = set()
+    
+    with torch.inference_mode():
+        memory_image = model.encode(image)
+        
+        for i in range(max_iter):
+            x = torch.tensor([current_text], dtype=torch.long).to(device)
+            output = model.decode(memory_image, x)
+            output = output[:, -1, :]
+            
+            for token_id in already_seen:
+                if output[0, token_id] < 0:
+                    output[0, token_id] *= penalty
+                else:
+                    output[0, token_id] /= penalty
+                    
+            output = output / temp
+            topk_vals, topk_idx = torch.topk(output, top_k, dim=-1)
+            probs = F.softmax(topk_vals, dim=-1)
+            idx = torch.multinomial(probs, 1)
+
+            idx = topk_idx.gather(-1, idx).item()
+            if idx == eos_id:
+                break
+
+            generated.append(idx)
+            current_text.append(idx)
+            already_seen.add(idx)
+
+    return generated
+
+def generate_fast(model, image, bos_id, eos_id, device="cuda", max_iter=32):
+    """
+    Super-fast prediction using only argmax (no sampling).
+
+    Args:
+        model: SOCRATE model instance.
+        image (Tensor): Single image tensor [1, C, H, W].
+        bos_id (int): Begin-of-sequence token ID.
+        eos_id (int): End-of-sequence token ID.
+        device (str): Target device. Default: "cuda".
+        max_iter (int): Max number of tokens to generate. Default: 32.
+    """
+    model.eval()
+    current_text = [bos_id]
+    generated = []
+    
+    with torch.inference_mode():
+        memory_image = model.encode(image)
+        for _ in range(max_iter):
+            x = torch.tensor([current_text], dtype=torch.long, device=device)
+            output = model.decode(memory_image, x)
+            logits = output[:, -1, :]
+            idx = logits.argmax(dim=-1).item()
+            
+            if idx == eos_id:
+                break
+                
+            generated.append(idx)
+            current_text.append(idx)
+            
+    return generated
+
+def beam_search(model, image, bos_id, eos_id, device="cuda", beam_width=4, max_iter=64):
+    """
+    Beam search decoding.
+
+    Args:
+        model: SOCRATE model instance.
+        image (Tensor): Single image tensor [1, C, H, W].
+        bos_id (int): Begin-of-sequence token ID.
+        eos_id (int): End-of-sequence token ID.
+        device (str): Target device. Default: "cuda".
+        beam_width (int): Number of beams. Default: 4.
+        max_iter (int): Max tokens per beam. Default: 64.
+    
+    Note: Full beam search is coming soon. Currently uses generate_fast as a fallback.
+    """
+    print("WARNING: Beam search is not fully implemented yet. Using generate_fast as a fallback.")
+    return generate_fast(model, image, bos_id, eos_id, device, max_iter=max_iter)
+
+def extract_crops_from_image(image_path, doctr_model=None):
+    """
+    Extracts words (crops) using doctr and sorts them
+    correctly from top-to-bottom and left-to-right.
+    """
+    if doctr_model is None:
+        from doctr.models import detection_predictor
+        doctr_model = detection_predictor(arch="db_resnet50", pretrained=True)
+        
+    from doctr.io import DocumentFile
+    
+    doc = DocumentFile.from_images(image_path)
+    result = doctr_model(doc)
+    
+    boxes = result[0]["words"]
+    image = cv2.imread(image_path)
+    if image is None:
+        raise ValueError(f"Could not load image from {image_path}")
+        
+    H, W = image.shape[:2]
+
+    # Sort by lines a la SOCRATE
+    boxes_info = []
+    for b in boxes:
+        xmin, ymin, xmax, ymax, score = b
+        cy = (ymin + ymax) / 2.0
+        h = ymax - ymin
+        boxes_info.append({'box': b, 'cy': cy, 'h': h, 'x': xmin})
+
+    boxes_info.sort(key=lambda item: item['cy'])
+
+    lines = []
+    current_line = []
+
+    for b in boxes_info:
+        if not current_line:
+            current_line.append(b)
+        else:
+            tolerance = current_line[0]['h'] * 0.5
+            if abs(b['cy'] - current_line[0]['cy']) < tolerance:
+                current_line.append(b)
+            else:
+                lines.append(current_line)
+                current_line = [b]
+    if current_line:
+        lines.append(current_line)
+
+    sorted_boxes = []
+    for line in lines:
+        line.sort(key=lambda item: item['x'])
+        for item in line:
+            sorted_boxes.append(item['box'])
+
+    crops = []
+    for b in sorted_boxes:
+        xmin, ymin, xmax, ymax, score = b
+        x1 = int(xmin * W)
+        y1 = int(ymin * H)
+        x2 = int(xmax * W)
+        y2 = int(ymax * H)
+
+        crop = image[y1:y2, x1:x2]
+        h, w = crop.shape[:2]
+
+        if h == 0 or w == 0:
+            continue
+            
+        crops.append(crop)
+        
+    return crops
+
+def predict(model, tokenizer, image_paths, wpb=16, function="generate_fast", doctr_model=None, bos_id=None, eos_id=None, device="cuda",
+            # generate() params
+            temp=None, max_iter=None, penalty=None, top_k=None,
+            # generate_fast() params
+            fast_max_iter=None,
+            # beam_search() params
+            beam_width=None, beam_max_iter=None):
+    """
+    The main prediction function of the library. 
+    Takes images and returns the text read from them.
+    
+    Inference parameters (temp, max_iter, penalty, top_k, fast_max_iter, beam_width, beam_max_iter)
+    can be set here directly, OR they will be read from model.sx_config if you created the model via sx.init(config=...).
+    
+    Args:
+        model: SOCRATE model instance.
+        tokenizer: SocrateXTokenizer instance.
+        image_paths (str | List[str]): Path(s) to the image(s).
+        wpb (int): Words per batch. Default: 16.
+        function (str | callable): 'generate', 'generate_fast', 'beam_search', or a custom callable.
+        doctr_model: Pre-loaded doctr model (avoids re-loading on each call).
+        bos_id (int): Override BOS token ID.
+        eos_id (int): Override EOS token ID.
+        device (str): Target device. Default: "cuda".
+        temp (float): Temperature for generate(). Default: from config or 0.5.
+        max_iter (int): Max tokens for generate(). Default: from config or 64.
+        penalty (float): Repetition penalty for generate(). Default: from config or 1.15.
+        top_k (int): Top-k for generate(). Default: from config or 5.
+        fast_max_iter (int): Max tokens for generate_fast(). Default: from config or 32.
+        beam_width (int): Number of beams for beam_search(). Default: from config or 4.
+        beam_max_iter (int): Max tokens for beam_search(). Default: from config or 64.
+    """
+    model.eval()
+
+    # Pull inference defaults from sx_config if they were set
+    sx_cfg = getattr(model, "sx_config", None)
+
+    _temp        = temp        if temp        is not None else (sx_cfg.temp        if sx_cfg else 0.5)
+    _max_iter    = max_iter    if max_iter    is not None else (sx_cfg.max_iter    if sx_cfg else 64)
+    _penalty     = penalty     if penalty     is not None else (sx_cfg.penalty     if sx_cfg else 1.15)
+    _top_k       = top_k       if top_k       is not None else (sx_cfg.top_k       if sx_cfg else 5)
+    _fast_max    = fast_max_iter if fast_max_iter is not None else (sx_cfg.fast_max_iter if sx_cfg else 32)
+    _beam_width  = beam_width  if beam_width  is not None else (sx_cfg.beam_width  if sx_cfg else 4)
+    _beam_max    = beam_max_iter if beam_max_iter is not None else (sx_cfg.beam_max_iter if sx_cfg else 64)
+    
+    # Resolve tokens (default to tokenizer if not provided)
+    if bos_id is None:
+        bos_id = tokenizer.token_to_id("<bos>")
+    if eos_id is None:
+        eos_id = tokenizer.token_to_id("<eos>")
+        
+    results = {}
+    
+    if isinstance(image_paths, str):
+        image_paths = [image_paths]
+        
+    for image_path in image_paths:
+        crops = extract_crops_from_image(image_path, doctr_model)
+        if not crops:
+            results[image_path] = ""
+            continue
+            
+        dataset = Makeset(images=crops)
+        dataloader = DataLoader(dataset, batch_size=wpb, shuffle=False, collate_fn=dataset.collate_fn)
+        
+        doc_text = []
+        for batch in dataloader:
+            batch = batch.to(device)
+            for img in batch:
+                img = img.unsqueeze(0) # [1, C, H, W]
+                
+                if function == "generate":
+                    pred_ids = generate(model, img, bos_id=bos_id, eos_id=eos_id, device=device,
+                                        temp=_temp, max_iter=_max_iter, penalty=_penalty, top_k=_top_k)
+                elif function == "generate_fast":
+                    pred_ids = generate_fast(model, img, bos_id=bos_id, eos_id=eos_id, device=device,
+                                             max_iter=_fast_max)
+                elif function == "beam_search":
+                    pred_ids = beam_search(model, img, bos_id=bos_id, eos_id=eos_id, device=device,
+                                           beam_width=_beam_width, max_iter=_beam_max)
+                else:
+                    if callable(function):
+                        pred_ids = function(model, img, bos_id, eos_id, device)
+                    else:
+                        raise ValueError(f"Unknown function: {function}")
+                
+                text = tokenizer.decode(pred_ids)
+                doc_text.append(text)
+                
+        results[image_path] = " ".join(doc_text)
+        
+    return results