Create vision_llm_text_extraction.py

new_templates/vision_llm_text_extraction.py

@@ -0,0 +1,180 @@
+def extract_text_from_images_deployable():
+    """
+    Deployable watsonx.ai function that extracts text from multiple images/PDFs using foundation models.
+    
+    Expected input payload:
+    {
+        "input_data": [{
+            "values": [["<image_url_1>", "<image_url_2>", ...], ["<optional_extraction_prompt>"]]
+        }]
+    }
+    
+    Returns:
+    {
+        "predictions": [{
+            "fields": ["extracted_texts"],
+            "values": [[["<extracted_text_1>", "<extracted_text_2>", ...]]]
+        }]
+    }
+    """
+    
+    import mimetypes
+    import base64
+    import requests
+    from urllib.parse import urlparse
+    import fitz
+    from ibm_watsonx_ai import APIClient, Credentials
+    from ibm_watsonx_ai.foundation_models import ModelInference
+    
+    # Initialize watsonx client (these should be set as environment variables)
+    import os
+    WX_URL = os.getenv('WX_URL', "")
+    WX_APIKEY = os.getenv('WX_APIKEY', "")
+    PROJECT_ID = os.getenv('PROJECT_ID', "")
+    CHAT_MODEL = os.getenv('CHAT_MODEL', 'mistralai/mistral-medium-2505')
+    
+    DEFAULT_EXTRACTION_PROMPT = '''Extract all text within the image in a markdown form as close as possible to the original, free of any additional outputs that are not in the text, including descriptions of the element, comments about making outputs, etc.'''
+    
+    wx_credentials = Credentials(
+        url=WX_URL,
+        api_key=WX_APIKEY
+    )
+    client = APIClient(credentials=wx_credentials, project_id=PROJECT_ID)
+    
+    def create_data_url(source, filename=None):
+        """Create data URL from bytes, file path, or URL. Returns list for PDFs."""
+        
+        if isinstance(source, str) and source.startswith(('http://', 'https://')):
+            content = requests.get(source).content
+            filename = filename or urlparse(source).path.split('/')[-1] or 'file'
+        elif isinstance(source, str):
+            with open(source, 'rb') as f:
+                content = f.read()
+            filename = filename or source
+        else: 
+            content = source
+            if not filename:
+                raise ValueError("filename required for bytes input")
+        
+        mime_type = mimetypes.guess_type(filename)[0] or 'application/octet-stream'
+        
+        if mime_type == 'application/pdf':
+            doc = fitz.open(stream=content, filetype="pdf")
+            result = []
+            for page in doc:
+                pix = page.get_pixmap(matrix=fitz.Matrix(1.5, 1.5))
+                img_data = pix.tobytes("png")
+                encoded = base64.b64encode(img_data).decode('utf-8')
+                result.append(f"data:image/png;base64,{encoded}")
+            doc.close()
+            return result
+        
+        encoded = base64.b64encode(content).decode('utf-8')
+        return f"data:{mime_type};base64,{encoded}"
+    
+    def score(payload):
+        """
+        Score function called for each prediction request.
+        
+        Args:
+            payload: Input payload containing list of image URLs/paths and optional extraction prompt
+            
+        Returns:
+            Dictionary with predictions containing list of extracted texts
+        """
+        try:
+            # Extract input data from payload
+            input_values = payload.get("input_data")[0].get("values")
+            image_urls = input_values[0]  # List of URLs
+            extraction_prompt = input_values[1] if len(input_values) > 1 else DEFAULT_EXTRACTION_PROMPT
+            
+            # Model parameters
+            params = {
+                "temperature": 1.0,
+                "max_tokens": 6553,
+                "top_p": 1.0,
+                "stop": [
+                    "</s>",
+                    "<|end_of_text|>"
+                ]
+            }
+            
+            extracted_texts = []
+            
+            # Process each image URL
+            for image_url in image_urls:
+                # Convert image to data URL
+                image_data_url = create_data_url(image_url)
+                
+                # Handle PDF case (multiple pages)
+                if isinstance(image_data_url, list):
+                    all_extracted_text = []
+                    for page_num, page_url in enumerate(image_data_url):
+                        messages = [
+                            {
+                                "role": "user",
+                                "content": [
+                                    {
+                                        "type": "text",
+                                        "text": f"Page {page_num + 1}:\n{extraction_prompt}"
+                                    },
+                                    {
+                                        "type": "image_url",
+                                        "image_url": {
+                                            "url": page_url,
+                                        }
+                                    }
+                                ]
+                            }
+                        ]
+                        
+                        chat_model = ModelInference(api_client=client, model_id=CHAT_MODEL, params=params)
+                        model_response = chat_model.chat(messages=messages)
+                        page_text = model_response["choices"][0]["message"]["content"]
+                        all_extracted_text.append(f"## Page {page_num + 1}\n\n{page_text}")
+                    
+                    extracted_text = "\n\n".join(all_extracted_text)
+                else:
+                    # Single image case
+                    messages = [
+                        {
+                            "role": "user",
+                            "content": [
+                                {
+                                    "type": "text",
+                                    "text": extraction_prompt
+                                },
+                                {
+                                    "type": "image_url",
+                                    "image_url": {
+                                        "url": image_data_url,
+                                    }
+                                }
+                            ]
+                        }
+                    ]
+                    
+                    chat_model = ModelInference(api_client=client, model_id=CHAT_MODEL, params=params)
+                    model_response = chat_model.chat(messages=messages)
+                    extracted_text = model_response["choices"][0]["message"]["content"]
+                
+                extracted_texts.append(extracted_text)
+            
+            # Return in required format
+            return {
+                'predictions': [{
+                    'fields': ['extracted_texts'],
+                    'values': [extracted_texts]
+                }]
+            }
+            
+        except Exception as e:
+            # Return error in predictions format
+            return {
+                'predictions': [{
+                    'fields': ['extracted_texts', 'error'],
+                    'values': [[], str(e)]
+                }]
+            }
+    
+    return score