Rename components/llm_utils.py to components/lllm_ocr.py

components/lllm_ocr.py

@@ -0,0 +1,37 @@
+from transformers import MllamaForConditionalGeneration, AutoProcessor
+from PIL import Image
+import torch
+
+# Model and processor init
+ckpt = "unsloth/Llama-3.2-11B-Vision-Instruct"
+model = MllamaForConditionalGeneration.from_pretrained(
+    ckpt,
+    torch_dtype=torch.bfloat16
+).to("cuda")
+processor = AutoProcessor.from_pretrained(ckpt)
+
+def extract_text_llm(image_path):
+    image = Image.open(image_path).convert("RGB")
+
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Extract handwritten text from the image and output only the extracted text without any additional description or commentary in output"},
+                {"type": "image"}
+            ]
+        }
+    ]
+
+    texts = processor.apply_chat_template(messages, add_generation_prompt=True)
+    inputs = processor(text=texts, images=[image], return_tensors="pt").to("cuda")
+
+    outputs = model.generate(**inputs, max_new_tokens=250)
+    result = processor.decode(outputs[0], skip_special_tokens=True)
+
+    # Clean result
+    if "assistant" in result.lower():
+        result = result[result.lower().find("assistant") + len("assistant"):].strip()
+    result = result.replace("user", "").strip()
+
+    return result