pembot 0.0.3__py2.py3-none-any.whl → 0.0.4__py2.py3-none-any.whl

pembot/AnyToText/convertor.py

@@ -7,6 +7,7 @@ import json
 import pandas as pd
 from typing import Literal, Union, Dict, Any, List
 import tempfile
+from datetime import datetime, date
 
 
 PandasReadEngineType = Literal['xlrd', 'openpyxl', 'odf', 'pyxlsb', 'calamine', None]
@@ -30,19 +31,19 @@ EXCEL_FILE_TYPES= [
 class Convertor():
 
 
-    def __init__(self, myfile: Path | None, output_dir: Path | None, file_bytes: bytes | None, suffix: str | None, file_type: str | None):
+    def __init__(self, myfile: Path | None= None, output_dir: Path | None= None, file_bytes: bytes | None= None, suffix: str | None= None, file_type: str | None= None):
 
         self.output= ""
 
         # file_type can be pdf, excel, etc.
-        if output_dir is None and file_bytes is not None and suffix is not None and myfile is None:
+        if output_dir is None and myfile is None and file_bytes is not None and suffix is not None:
             with tempfile.TemporaryDirectory() as dp:
                 with tempfile.NamedTemporaryFile(suffix= suffix, mode= 'wb') as fp:
                     fp.write(file_bytes)
                     myfile= Path(fp.name)
                     output_dir= Path(dp)
                     if file_type == 'pdf':
-                        extractor= MarkdownPDFExtractor(str(myfile), output_path= str(self.output_dir), page_delimiter= "-- NEXT PAGE --")
+                        extractor= MarkdownPDFExtractor(str(myfile), output_path= str(output_dir), page_delimiter= "-- NEXT PAGE --")
                         extractor.extract()
                         with open(output_dir / (myfile.stem + '.md')) as output_file:
                             self.output= output_file.read()
@@ -75,162 +76,235 @@ class Convertor():
             else:
                 print(mt)
 
-
-
     def convert_file_to_json(
-            self,
-            sheet_to_convert: Union[str, int, None] = None,  # Relevant for Excel/ODS
-            orient: Literal['dict', 'list', 'series', 'split', 'records', 'index'] = 'records', # Corrected type hint
-            date_format: Union[str, None] = 'iso',  # 'iso', 'epoch', or None
-            csv_encoding: str = 'utf-8', # For reading CSV files
-            excel_ods_engine: PandasReadEngineType = None # For Excel/ODS, e.g., 'openpyxl', 'xlrd', 'odf'
+        self,
+        sheet_to_convert: Union[str, int, None] = None,  # Relevant for Excel/ODS
+        orient: Literal['dict', 'list', 'series', 'split', 'records', 'index'] = 'records', # Corrected type hint
+        date_format: Union[str, None] = 'iso',  # 'iso', 'epoch', or None
+        csv_encoding: str = 'utf-8', # For reading CSV files
+        excel_ods_engine: PandasReadEngineType = None # For Excel/ODS, e.g., 'openpyxl', 'xlrd', 'odf'
         ) -> bool:
-            """
-            Converts an Excel, ODS, or CSV file (or a specific Excel/ODS sheet)
-            into an equivalent JSON format.
-
-            Args:
-                sheet_to_convert (str | int | None, optional):
-                    - For Excel/ODS:
-                        - If None (default): Converts all sheets. The JSON output will be a
-                          dictionary where keys are sheet names and values are the JSON
-                          representation of each sheet.
-                        - If str: Name of the specific sheet to convert.
-                        - If int: Index of the specific sheet to convert (0-based).
-                        If a specific sheet is requested, the JSON output will directly be
-                        the representation of that sheet.
-                    - For CSV: This parameter is ignored. The entire CSV is processed.
-                orient (str, optional): Pandas DataFrame.to_dict() orientation for each sheet/CSV.
-                    Default: 'records'. See pandas.DataFrame.to_dict() documentation.
-                date_format (str | None, optional): Format for datetime objects.
-                    - 'iso' (default): ISO8601 format (e.g., '2023-10-27T10:30:00').
-                    - 'epoch': Milliseconds since epoch.
-                    - None: Pandas default (often Timestamps). 'iso' is generally safer for JSON.
-                csv_encoding (str, optional): Encoding for reading CSV files. Default is 'utf-8'.
-                excel_ods_engine (str | None, optional): Pandas engine for reading Excel or ODS files.
-                    - For Excel: 'openpyxl' (for .xlsx), 'xlrd' (for .xls).
-                    - For ODS: 'odf' (requires 'odfpy' library).
-                    If None, pandas auto-detects based on file extension and installed libraries.
-
-            Returns:
-                bool: True if conversion was successful, False otherwise.
-            """
-            input_filepath = self.input_filepath
-            json_filepath = self.json_filepath
-
-            try:
-
-                if not input_filepath.exists():
-                    print(f"Error: Input file not found at {input_filepath}")
-                    return False
+        """
+        Converts an Excel, ODS, or CSV file (or a specific Excel/ODS sheet)
+        into an equivalent JSON format.
+
+        Args:
+            sheet_to_convert (str | int | None, optional):
+                - For Excel/ODS:
+                    - If None (default): Converts all sheets. The JSON output will be a
+                      dictionary where keys are sheet names and values are the JSON
+                      representation of each sheet.
+                    - If str: Name of the specific sheet to convert.
+                    - If int: Index of the specific sheet to convert (0-based).
+                    If a specific sheet is requested, the JSON output will directly be
+                    the representation of that sheet.
+                - For CSV: This parameter is ignored. The entire CSV is processed.
+            orient (str, optional): Pandas DataFrame.to_dict() orientation for each sheet/CSV.
+                Default: 'records'. See pandas.DataFrame.to_dict() documentation.
+            date_format (str | None, optional): Format for datetime objects.
+                - 'iso' (default): ISO8601 format (e.g., '2023-10-27T10:30:00').
+                - 'epoch': Milliseconds since epoch.
+                - None: Pandas default (often Timestamps). 'iso' is generally safer for JSON.
+            csv_encoding (str, optional): Encoding for reading CSV files. Default is 'utf-8'.
+            excel_ods_engine (str | None, optional): Pandas engine for reading Excel or ODS files.
+                - For Excel: 'openpyxl' (for .xlsx), 'xlrd' (for .xls).
+                - For ODS: 'odf' (requires 'odfpy' library).
+                If None, pandas auto-detects based on file extension and installed libraries.
+
+        Returns:
+            bool: True if conversion was successful, False otherwise.
+        """
+
+        input_filepath = self.input_filepath
+        json_filepath = self.json_filepath
+
+        try:
+
+            if not input_filepath.exists():
+                print(f"Error: Input file not found at {input_filepath}")
+                return False
 
-                # Ensure output directory exists
-                json_filepath.parent.mkdir(parents=True, exist_ok=True)
+            # Ensure output directory exists
+            json_filepath.parent.mkdir(parents=True, exist_ok=True)
 
-                file_suffix = input_filepath.suffix.lower()
-                output_data_final: Union[Dict[str, Any], List[Dict[str, Any]]] = {}
+            file_suffix = input_filepath.suffix.lower()
+            output_data_final: Union[Dict[str, Any], List[Dict[str, Any]]] = {}
 
-                dataframes_to_process: list[tuple[pd.DataFrame, str | None]] = []
+            dataframes_to_process: list[tuple[pd.DataFrame, str | None]] = []
 
-                current_engine: PandasReadEngineType = excel_ods_engine
+            current_engine: PandasReadEngineType = excel_ods_engine
 
-                if file_suffix == '.csv':
-                    if sheet_to_convert is not None:
-                        print(f"Info: 'sheet_to_convert' parameter ('{sheet_to_convert}') is ignored for CSV file '{input_filepath.name}'. Processing entire CSV.")
-                    try:
-                        df = pd.read_csv(input_filepath, encoding=csv_encoding)
-                        dataframes_to_process.append((df, None))
-                    except Exception as e:
-                        print(f"Error reading CSV file '{input_filepath.name}': {e}")
-                        return False
-
-                elif file_suffix in ['.xls', '.xlsx', '.ods']:
-                    try:
-                        if file_suffix == '.ods':
-                            if current_engine is None:
-                                current_engine = 'odf'
-                            elif current_engine != 'odf':
-                                print(f"Warning: Specified engine '{current_engine}' may not be optimal for ODS. Forcing 'odf'.")
-                                current_engine = 'odf'
-
-                        if sheet_to_convert is not None:
-                            df = pd.read_excel(input_filepath, sheet_name=sheet_to_convert, engine=current_engine)
-                            dataframes_to_process.append((df, None))
-
-                        else:
-                            excel_file = pd.ExcelFile(input_filepath, engine=current_engine)
-                            if not excel_file.sheet_names:
-                                print(f"Warning: File '{input_filepath.name}' contains no sheets.")
-                            for sheet_name in excel_file.sheet_names:
-                                df = excel_file.parse(sheet_name) # engine is inherited
-                                dataframes_to_process.append((df, sheet_name))
-                    except ImportError as ie:
-                        if 'odfpy' in str(ie).lower() and file_suffix == '.ods':
-                            print(f"Error reading ODS file '{input_filepath.name}': The 'odfpy' library is required. Please install it using 'pip install odfpy'.")
-                        elif 'xlrd' in str(ie).lower() and file_suffix == '.xls':
-                            print(f"Error reading .xls file '{input_filepath.name}': The 'xlrd' library might be required. Please install it using 'pip install xlrd'.")
-                        elif 'openpyxl' in str(ie).lower() and file_suffix == '.xlsx':
-                            print(f"Error reading .xlsx file '{input_filepath.name}': The 'openpyxl' library might be required. Please install it using 'pip install openpyxl'.")
-                        else:
-                            print(f"ImportError reading file '{input_filepath.name}': {ie}")
-                        return False
-                    except Exception as e:
-                        print(f"Error reading Excel/ODS file '{input_filepath.name}': {e}")
-                        return False
-                else:
-                    print(f"Error: Unsupported file type: '{file_suffix}'. Please provide a CSV, XLS, XLSX, or ODS file.")
+            if file_suffix == '.csv':
+                if sheet_to_convert is not None:
+                    print(f"Info: 'sheet_to_convert' parameter ('{sheet_to_convert}') is ignored for CSV file '{input_filepath.name}'. Processing entire CSV.")
+                try:
+                    df = pd.read_csv(input_filepath, encoding=csv_encoding)
+                    dataframes_to_process.append((df, None))
+                except Exception as e:
+                    print(f"Error reading CSV file '{input_filepath.name}': {e}")
                     return False
 
-                if not dataframes_to_process and file_suffix in ['.xls', '.xlsx', '.ods'] and sheet_to_convert is None:
-                     print(f"Info: No dataframes were loaded from '{input_filepath.name}'. Output JSON will be empty if processing all sheets from an empty file.")
-                elif not dataframes_to_process and not (file_suffix in ['.xls', '.xlsx', '.ods'] and sheet_to_convert is None):
-                     pass
+            elif file_suffix in ['.xls', '.xlsx', '.ods']:
+                try:
+                    if file_suffix == '.ods':
+                        if current_engine is None:
+                            current_engine = 'odf'
+                        elif current_engine != 'odf':
+                            print(f"Warning: Specified engine '{current_engine}' may not be optimal for ODS. Forcing 'odf'.")
+                            current_engine = 'odf'
 
+                    if sheet_to_convert is not None:
+                        df = pd.read_excel(input_filepath, sheet_name=sheet_to_convert, engine=current_engine)
+                        dataframes_to_process.append((df, None))
 
-                is_direct_output = len(dataframes_to_process) == 1 and dataframes_to_process[0][1] is None
-                temp_processed_data: Dict[str, Any] = {}
-
-                for df_original, name_key in dataframes_to_process:
-                    df = df_original.copy()
-
-                    if date_format:
-                        for col_name in df.select_dtypes(include=['datetime64[ns]', 'datetime', 'datetimetz']).columns:
-                            try:
-                                if date_format == 'iso':
-                                    df[col_name] = df[col_name].apply(lambda x: x.isoformat() if pd.notnull(x) and hasattr(x, 'isoformat') else None)
-                                elif date_format == 'epoch':
-                                    df[col_name] = df[col_name].apply(lambda x: int(x.timestamp() * 1000) if pd.notnull(x) and hasattr(x, 'timestamp') else None)
-                            except Exception as e_date:
-                                print(f"Warning: Could not fully convert date column '{col_name}' in '{name_key or input_filepath.name}' using format '{date_format}'. Error: {e_date}. Problematic values might be None.")
-
-                    df = df.astype(object).where(pd.notnull(df), None)
-                    current_json_segment = df.to_dict(orient=orient)
-
-                    if is_direct_output:
-                        output_data_final = current_json_segment
-                        break
                     else:
-                        if name_key is not None:
-                            temp_processed_data[name_key] = current_json_segment
-
-                if not is_direct_output:
-                    output_data_final = temp_processed_data
-
-                with open(json_filepath, 'w', encoding='utf-8') as f:
-                    json.dump(output_data_final, f, indent=4, ensure_ascii=False)
+                        excel_file = pd.ExcelFile(input_filepath, engine=current_engine)
+                        if not excel_file.sheet_names:
+                            print(f"Warning: File '{input_filepath.name}' contains no sheets.")
+                        for sheet_name in excel_file.sheet_names:
+                            df = excel_file.parse(sheet_name) # engine is inherited
+                            dataframes_to_process.append((df, sheet_name))
+                except ImportError as ie:
+                    if 'odfpy' in str(ie).lower() and file_suffix == '.ods':
+                        print(f"Error reading ODS file '{input_filepath.name}': The 'odfpy' library is required. Please install it using 'pip install odfpy'.")
+                    elif 'xlrd' in str(ie).lower() and file_suffix == '.xls':
+                        print(f"Error reading .xls file '{input_filepath.name}': The 'xlrd' library might be required. Please install it using 'pip install xlrd'.")
+                    elif 'openpyxl' in str(ie).lower() and file_suffix == '.xlsx':
+                        print(f"Error reading .xlsx file '{input_filepath.name}': The 'openpyxl' library might be required. Please install it using 'pip install openpyxl'.")
+                    else:
+                        print(f"ImportError reading file '{input_filepath.name}': {ie}")
+                    return False
+                except Exception as e:
+                    print(f"Error reading Excel/ODS file '{input_filepath.name}': {e}")
+                    return False
+            else:
+                print(f"Error: Unsupported file type: '{file_suffix}'. Please provide a CSV, XLS, XLSX, or ODS file.")
+                return False
 
-                print(f"Successfully converted '{input_filepath.name}' to '{json_filepath.name}'")
-                return True
+            if not dataframes_to_process and file_suffix in ['.xls', '.xlsx', '.ods'] and sheet_to_convert is None:
+                 print(f"Info: No dataframes were loaded from '{input_filepath.name}'. Output JSON will be empty if processing all sheets from an empty file.")
+            elif not dataframes_to_process and not (file_suffix in ['.xls', '.xlsx', '.ods'] and sheet_to_convert is None):
+                 pass
+
+            is_direct_output = len(dataframes_to_process) == 1 and dataframes_to_process[0][1] is None
+            temp_processed_data: Dict[str, Any] = {}
+
+            for df_original, name_key in dataframes_to_process:
+                df = df_original.copy()
+
+                # Handle datetime columns with improved detection and conversion
+                if date_format:
+                    # Check for datetime columns using multiple approaches
+                    datetime_columns = []
+
+                    # Method 1: Use pandas dtype detection
+                    datetime_columns.extend(df.select_dtypes(include=['datetime64[ns]', 'datetime', 'datetimetz']).columns.tolist())
+
+                    # Method 2: Check for datetime objects in each column
+                    for col in df.columns:
+                        if col not in datetime_columns:
+                            # Sample a few non-null values to check type
+                            sample_values = df[col].dropna().head(10)
+                            if len(sample_values) > 0:
+                                for val in sample_values:
+                                    if isinstance(val, (datetime, date, pd.Timestamp)):
+                                        datetime_columns.append(col)
+                                        break
+
+                    # Convert datetime columns
+                    for col_name in datetime_columns:
+                        try:
+                            if date_format == 'iso':
+                                df[col_name] = df[col_name].apply(lambda x: self._convert_to_iso(x))
+                            elif date_format == 'epoch':
+                                df[col_name] = df[col_name].apply(lambda x: self._convert_to_epoch(x))
+                        except Exception as e_date:
+                            print(f"Warning: Could not fully convert date column '{col_name}' in '{name_key or input_filepath.name}' using format '{date_format}'. Error: {e_date}")
+
+                # Replace NaN values with None for JSON compatibility
+                df = df.astype(object).where(pd.notnull(df), None)
+
+                # Final safety check: convert any remaining datetime objects
+                for col in df.columns:
+                    df[col] = df[col].apply(lambda x: self._safe_datetime_convert(x, date_format))
+
+                current_json_segment = df.to_dict(orient=orient)
+
+                if is_direct_output:
+                    output_data_final = current_json_segment
+                    break
+                else:
+                    if name_key is not None:
+                        temp_processed_data[name_key] = current_json_segment
+
+            if not is_direct_output:
+                output_data_final = temp_processed_data
+
+            with open(json_filepath, 'w', encoding='utf-8') as f:
+                json.dump(output_data_final, f, indent=4, ensure_ascii=False)
+
+            print(f"Successfully converted '{input_filepath.name}' to '{json_filepath.name}'")
+            return True
+
+        except FileNotFoundError:
+            print(f"Error: Input file not found at {input_filepath.name}")
+            return False
+        except ValueError as ve:
+            print(f"ValueError during conversion of '{input_filepath.name}': {ve}")
+            return False
+        except Exception as e:
+            print(f"An unexpected error occurred during conversion of '{input_filepath.name}': {e}")
+            return False
+
+    def _convert_to_iso(self, value):
+        """Convert datetime-like objects to ISO format string."""
+        if pd.isnull(value) or value is None:
+            return None
+
+        try:
+            if isinstance(value, str):
+                return value  # Already a string
+            elif hasattr(value, 'isoformat'):
+                return value.isoformat()
+            elif isinstance(value, pd.Timestamp):
+                return value.isoformat()
+            else:
+                return str(value)
+        except:
+            return str(value) if value is not None else None
+
+    def _convert_to_epoch(self, value):
+        """Convert datetime-like objects to epoch milliseconds."""
+        if pd.isnull(value) or value is None:
+            return None
+
+        try:
+            if isinstance(value, (int, float)):
+                return int(value)  # Assume already epoch
+            elif hasattr(value, 'timestamp'):
+                return int(value.timestamp() * 1000)
+            elif isinstance(value, pd.Timestamp):
+                return int(value.timestamp() * 1000)
+            else:
+                return str(value)
+        except:
+            return str(value) if value is not None else None
+
+    def _safe_datetime_convert(self, value, date_format):
+        """Final safety conversion for any remaining datetime objects."""
+        if pd.isnull(value) or value is None:
+            return None
+
+        # If it's a datetime-like object, convert it
+        if isinstance(value, (datetime, date, pd.Timestamp)):
+            if date_format == 'iso':
+                return self._convert_to_iso(value)
+            elif date_format == 'epoch':
+                return self._convert_to_epoch(value)
+            else:
+                return str(value)
 
-            except FileNotFoundError:
-                print(f"Error: Input file not found at {input_filepath.name}")
-                return False
-            except ValueError as ve:
-                print(f"ValueError during conversion of '{input_filepath.name}': {ve}")
-                return False
-            except Exception as e:
-                print(f"An unexpected error occurred during conversion of '{input_filepath.name}': {e}")
-                return False
+        return value
 
 
 def chunk_text(text, chunk_size=500, overlap_size=50):
@@ -257,4 +331,34 @@ def chunk_text(text, chunk_size=500, overlap_size=50):
     return chunks
 
 if __name__ == '__main__':
-    print("do you want a rice bag?")
+    print("Test Run Start:")
+    try:
+        # print("Test 1: scaned pdf page, bytes")
+        # with open("/home/cyto/Documents/scanned.pdf", "rb") as imgpdf:
+        #     conv= Convertor(file_bytes= imgpdf.read(), suffix= ".pdf", file_type= "pdf")
+        #     print(conv.output)
+
+        # print("Test 2: JD pdf, bytes")
+        # with open("/home/cyto/dev/pembotdir/jds/PM Trainee.pdf", "rb") as imgpdf:
+        #     conv= Convertor(file_bytes= imgpdf.read(), suffix= ".pdf", file_type= "pdf")
+        #     print(conv.output)
+
+        # print("Test 3: excel schedule, bytes")
+        # with open("/home/cyto/Downloads/Assignment schedule.xlsx", "rb") as imgpdf:
+        #     conv= Convertor(file_bytes= imgpdf.read(), suffix= ".xlsx", file_type= "excel")
+        #     print(conv.output)
+
+        # without bytes example:
+        print("Test 4: scanned pdf, path")
+        conv= Convertor(myfile= Path('/home/cyto/Documents/scanned.pdf'), output_dir= Path('/home/cyto/Documents'))
+        print(conv.output)
+
+        # print("Test 5: schedule excel, path")
+        # conv= Convertor(myfile= Path('/home/cyto/Downloads/Assignment schedule.xlsx'), output_dir= Path('/home/cyto/Downloads'))
+        # print(conv.output)
+    except FileNotFoundError as fe:
+        print("file not found, modify the driver code to get sample files to test:\n\n", fe)
+    except Exception as e:
+        print("unhandled: ", e)
+
+    print("Test Run End.")

pembot/__init__.py

@@ -1,6 +1,6 @@
 """
 A Python Package to convert PEM blog content to usseful information by leveraging LLMs
 """
-__version__ = '0.0.3'
+__version__ = '0.0.4'
 from .main import save_to_json_file, make_query
 __all__ = ["save_to_json_file", "make_query"]

pembot/config/config.yaml

@@ -2,4 +2,4 @@ OUTPUT_DIR: /home/cyto/dev/pembotdir
 PAGE_DELIMITER: ___________________________ NEXT PAGE ___________________________
 app:
   name: pembot
-  version: 0.0.3
+  version: 0.0.4

pembot/main.py

@@ -10,7 +10,11 @@ from pembot.query import rag_query_llm, remove_bs
 import os
 import json
 from pembot.utils.string_tools import make_it_an_id
-from schema.structure import required_fields
+import pickle
+from sys import argv
+
+required_fields_path= ""
+required_fields= None
 
 
 def make_query(required_fields: list[tuple[str, str, str, str]]):
@@ -67,8 +71,8 @@ def save_to_json_file(llm_output: str, filepath: Path):
     except Exception as e:
         print(f"An unexpected error occurred in save_to_json_file: {e}")
 
-def make_document_summarization_and_embeddings(db_client, llm_client, inference_client, docs_dir: Path, text_out_dir: Path, required_fields: list[tuple[str, str, str, str]], chunk_size: int = 600, embedding_model: str= 'nomic-embed-text:v1.5', llm_provider_name: PROVIDER_T= "novita", model_name= "deepseek-ai/DeepSeek-R1-0528-Qwen3-8B", embeddings_collection: str= "doc_chunks", index_name= "test_search"): 
-    # give required output fields 
+def make_document_summarization_and_embeddings(db_client, llm_client, inference_client, docs_dir: Path, text_out_dir: Path, required_fields: list[tuple[str, str, str, str]], chunk_size: int = 600, embedding_model: str= 'nomic-embed-text:v1.5', llm_provider_name: PROVIDER_T= "novita", model_name= "deepseek-ai/DeepSeek-R1-0528-Qwen3-8B", embeddings_collection: str= "doc_chunks", index_name= "test_search"):
+    # give required output fields
     # take the documents
     # convert to text
     # upload to chromadb
@@ -80,7 +84,7 @@ def make_document_summarization_and_embeddings(db_client, llm_client, inference_
         expected_json= text_out_dir / 'json' / (file_root + '.json')
         document_id= make_it_an_id(file_root)
 
-        if docfile.is_file and not (expected_json).exists(): 
+        if docfile.is_file and not (expected_json).exists():
 
             expected_markdown= text_out_dir / (file_root + '.md')
             if not (expected_markdown).exists():
@@ -161,6 +165,23 @@ if __name__ == "__main__":
     #     provider="Jina AI",
     #     api_key= JINA_API_KEY,
     # )
+    #
+
+    try:
+        if len(argv) > 1:
+            print(f"First argument: {argv[1]}")
+            required_fields_path= argv[1]
+            with open(required_fields_path, "rb") as rf:
+                required_fields= pickle.load(rf)
+    except Exception as e:
+        print("error while getting required_fields pickle. Please pickle it and put it in project directory to continue\n", e)
+
+    if required_fields is None:
+        print("couldnt load required fields. please provide path to pickle in command line argument")
+        exit()
+    else:
+        print(required_fields)
+
 
     inference_client= InferenceClient(
         provider="hf-inference",
@@ -178,7 +199,7 @@ if __name__ == "__main__":
     llm_provider_name: PROVIDER_T="nebius"
 
     # nerfed, but provided by hf serverless inference: BAAI/bge-small-en-v1.5
-    # Worth mentioning: 
+    # Worth mentioning:
     # jinaai/jina-embeddings-v3
     # BAAI/bge-base-en-v1.5
     # nomic-ai/nomic-embed-text-v1.5
@@ -203,6 +224,3 @@ if __name__ == "__main__":
 
     docs_collection= database["summary_docs"]
     upload_summaries(process_output_dir / 'json', docs_collection)
-
-
-

pembot/pdf2markdown/extract.py

@@ -1,11 +1,10 @@
-import fitz  # PyMuPDF
+import fitz
 import pdfplumber
 import re
 import yaml
 # import pytesseract
 import numpy as np
-from transformers import AutoTokenizer, AutoProcessor, AutoModelForImageTextToText
-# VisionEncoderDecoderModel, ViTImageProcessor,
+from transformers import AutoTokenizer, AutoProcessor, AutoModelForImageTextToText, VisionEncoderDecoderModel, ViTImageProcessor
 from typing import Literal, final
 import torch
 from PIL import Image
@@ -16,22 +15,11 @@ import warnings
 from pathlib import Path
 from abc import ABC, abstractmethod
 import argparse
-from PIL import Image
 import io
-from PIL import Image
-
-model_path = "nanonets/Nanonets-OCR-s"
-
-model = AutoModelForImageTextToText.from_pretrained(
-    model_path,
-    torch_dtype="auto",
-    device_map="auto",
-    attn_implementation="flash_attention_2"
-)
-model.eval()
+from google import genai
+from google.genai import types
+import mimetypes
 
-tokenizer = AutoTokenizer.from_pretrained(model_path)
-processor = AutoProcessor.from_pretrained(model_path)
 
 
 warnings.filterwarnings("ignore")
@@ -74,9 +62,31 @@ class MarkdownPDFExtractor(PDFExtractor):
 
     BULLET_POINTS = "•◦▪▫●○"
 
-    def __init__(self, pdf_path, output_path= config["OUTPUT_DIR"], page_delimiter= config["PAGE_DELIMITER"]):
+    def __init__(self, pdf_path, output_path= config["OUTPUT_DIR"], page_delimiter= config["PAGE_DELIMITER"], model_name: str | None= None):
         super().__init__(pdf_path)
 
+        if model_name is None:
+            self.MODEL_NAME= "gemini-2.5-flash"
+        else:
+            self.MODEL_NAME= model_name
+
+        if  "gemini" in self.MODEL_NAME:
+            self.gclient = genai.Client(api_key= os.getenv("GEMINI_API_KEY", ''))
+        else:
+            model_path = "nanonets/Nanonets-OCR-s"
+            self.model = AutoModelForImageTextToText.from_pretrained(
+                model_path,
+                torch_dtype="auto",
+                device_map="auto",
+                attn_implementation="flash_attention_2"
+            )
+            self.model.eval()
+            self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+            self.processor = AutoProcessor.from_pretrained(model_path)
+            self.setup_image_captioning()
+
+
+
         self.markdown_content= ""
         self.pdf_filename = Path(pdf_path).stem
         self.output_path= output_path
@@ -87,26 +97,26 @@ class MarkdownPDFExtractor(PDFExtractor):
         self.page_delimiter= page_delimiter
         Path(output_path).mkdir(parents=True, exist_ok=True)
 
-        # self.setup_image_captioning()
-
-    # def setup_image_captioning(self):
-    #     """Set up the image captioning model."""
-    #     try:
-    #         self.model = VisionEncoderDecoderModel.from_pretrained(
-    #             "nlpconnect/vit-gpt2-image-captioning"
-    #         )
-    #         self.feature_extractor = ViTImageProcessor.from_pretrained(
-    #             "nlpconnect/vit-gpt2-image-captioning"
-    #         )
-    #         self.tokenizer = AutoTokenizer.from_pretrained(
-    #             "nlpconnect/vit-gpt2-image-captioning"
-    #         )
-    #         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    #         self.model.to(self.device)
-    #         self.logger.info("Image captioning model set up successfully.")
-    #     except Exception as e:
-    #         self.logger.error(f"Error setting up image captioning model: {e}")
-    #         self.logger.exception(traceback.format_exc())
+
+
+    def setup_image_captioning(self):
+        """Set up the image captioning model."""
+        try:
+            self.model = VisionEncoderDecoderModel.from_pretrained(
+                "nlpconnect/vit-gpt2-image-captioning"
+            )
+            self.feature_extractor = ViTImageProcessor.from_pretrained(
+                "nlpconnect/vit-gpt2-image-captioning"
+            )
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                "nlpconnect/vit-gpt2-image-captioning"
+            )
+            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            self.model.to(self.device)
+            self.logger.info("Image captioning model set up successfully.")
+        except Exception as e:
+            self.logger.error(f"Error setting up image captioning model: {e}")
+            self.logger.exception(traceback.format_exc())
 
     def extract(self):
         try:
@@ -123,282 +133,197 @@ class MarkdownPDFExtractor(PDFExtractor):
             self.logger.exception(traceback.format_exc())
             return "", []
 
-    def extract_markdown_by_blocks(self):
-        """Main method to extract markdown from PDF."""
-        try:
-            doc = fitz.open(self.pdf_path)
-            markdown_content = ""
-            markdown_pages = []
-            tables = self.extract_tables()
-            table_index = 0
-            list_counter = 0
-            in_code_block = False
-            code_block_content = ""
-            code_block_lang = None
-            prev_line = ""
-
-            for page_num, page in enumerate(doc):
-                self.logger.info(f"Processing page {page_num + 1}")
-                page_content = ""
-                blocks = page.get_text("dict")["blocks"]
-                page_height = page.rect.height
-                links = self.extract_links(page)
-
-                if len(page.get_images()) > 0 and len(page.get_images()) <= 128:
-                    for block in blocks:
-                        if block["type"] == 0:  # Text
-                            page_content += self.process_text_block(
-                                block,
-                                page_height,
-                                links,
-                                list_counter,
-                                in_code_block,
-                                code_block_content,
-                                code_block_lang,
-                                prev_line,
-                            )
-                        elif block["type"] == 1:  # Image
-                            page_content += self.process_image_block(page, block)
-
-                else:
-                    for block in blocks:
-                        if block["type"] == 0:  # Text
-                            page_content += self.process_text_block(
-                                block,
-                                page_height,
-                                links,
-                                list_counter,
-                                in_code_block,
-                                code_block_content,
-                                code_block_lang,
-                                prev_line,
-                            )
-
-                # Insert tables at their approximate positions
-                while (
-                    table_index < len(tables)
-                    and tables[table_index]["page"] == page.number
-                ):
-                    page_content += (
-                        "\n\n"
-                        + self.table_to_markdown(tables[table_index]["content"])
-                        + "\n\n"
-                    )
-                    table_index += 1
-
-                markdown_pages.append(self.post_process_markdown(page_content))
-                markdown_content += page_content + config["PAGE_DELIMITER"]
 
-            markdown_content = self.post_process_markdown(markdown_content)
-            return markdown_content, markdown_pages
-        except Exception as e:
-            self.logger.error(f"Error extracting markdown: {e}")
-            self.logger.exception(traceback.format_exc())
-            return "", []
-
-
-    def ocr_page_with_nanonets_s(self, pil_image, model, processor, max_new_tokens: int | None = None):
+    def ocr_page_with_nanonets_s(self, pil_image, img_bytes, max_new_tokens: int | None = None):
         prompt = """Extract the text from the above document as if you were reading it naturally. Return the tables in html format. Return the equations in LaTeX representation. If there is an image in the document and image caption is not present, add a small description of the image inside the <img></img> tag; otherwise, add the image caption inside <img></img>. Watermarks should be wrapped in brackets. Ex: <watermark>OFFICIAL COPY</watermark>. Page numbers should be wrapped in brackets. Ex: <page_number>14</page_number> or <page_number>9/22</page_number>. Prefer using ☐ and ☑ for check boxes."""
         if max_new_tokens is None:
             max_new_tokens= 4096
 
-        # image = Image.open(image_path)
-        image = pil_image
-        messages = [
-            {"role": "system", "content": "You are a helpful assistant."},
-            {"role": "user", "content": [
-                {"type": "image", "image": image},
-                {"type": "text", "text": prompt},
-            ]},
-        ]
-        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        inputs = processor(text=[text], images=[image], padding=True, return_tensors="pt")
-        inputs = inputs.to(model.device)
+        if 'gemini' in self.MODEL_NAME:
+
+            image_format = pil_image.format
+            dummy_filename = f"dummy.{image_format.lower()}"
+            mime_type, _ = mimetypes.guess_type(dummy_filename)
+            response=  self.gclient.models.generate_content(
+                model= self.MODEL_NAME,
+                contents=[
+                    types.Part.from_bytes(
+                        data=img_bytes.getvalue(),
+                        mime_type= mime_type
+                    ),
+                    prompt
+                ]
+            )
+            # print("response :", response)
+            return response.text
+        else:
+            image = pil_image
+            messages = [
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": [
+                    {"type": "image", "image": image},
+                    {"type": "text", "text": prompt},
+                ]},
+            ]
+            text = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+            inputs = self.processor(text=[text], images=[image], padding=True, return_tensors="pt")
+            inputs = inputs.to(self.model.device)
 
-        output_ids = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False)
-        generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(inputs.input_ids, output_ids)]
+            output_ids = self.model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False)
+            generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(inputs.input_ids, output_ids)]
 
-        output_text = processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)
-        return output_text[0]
+            output_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)
+            return output_text[0]
 
 
 
     def extract_markdown(self):
-        """
-        Extracts all possible text content from a PDF page, concatenating it
-        from direct text blocks, OCR from embedded image blocks, and OCR from
-        full-page raster images (scanned pages).
+            """
+            Extracts all possible content from a PDF, prioritizing searchable text,
+            then OCR for embedded images, and finally full-page OCR for scanned pages.
+            Avoids redundant OCR where possible.
+
+            Returns:
+                tuple: A tuple containing:
+                    - str: The concatenated markdown content of all pages.
+                    - list: A list of strings, where each string is the comprehensive markdown
+                            for a corresponding page.
+            """
+            all_pages_markdown = []
+            full_document_markdown = [] # Changed to list of lines/blocks to handle insertions better
+
+            try:
+                doc = fitz.open(self.pdf_path)
+                self.logger.info(f"Opened PDF: {self.pdf_path}")
+
+                tables = self.extract_tables()
+                table_index = 0
+
+                # State variables for process_text_block that might need to persist across blocks
+                # Re-initialize for each new document, but allow state management within process_text_block for lines
+                list_counter = 0
+                in_code_block = False
+                code_block_content = ""
+                code_block_lang = None
+                prev_line = ""
+
+                for page_num, page in enumerate(doc):
+                    current_page_markdown_blocks = [] # Collect markdown blocks for the current page
+                    page_has_searchable_text = False
+                    page_has_embedded_images = False
+
+                    self.logger.info(f"\nProcessing page {page_num + 1}...")
+
+                    blocks = page.get_text('dict')['blocks']
+                    page_height = page.rect.height
+                    links = self.extract_links(page)
+
+                    # Phase 1: Process text blocks and embedded image blocks
+                    for block_num, block in enumerate(blocks):
+                        if block['type'] == 0:  # Text block
+                            page_has_searchable_text = True
+                            processed_text = self.process_text_block(
+                                block,
+                                page_height,
+                                links,
+                                list_counter,
+                                in_code_block,
+                                code_block_content,
+                                code_block_lang,
+                                prev_line,
+                            )
+                            if processed_text.strip():
+                                current_page_markdown_blocks.append(processed_text)
+
+                        elif block['type'] == 1:  # Image block
+                            page_has_embedded_images = True
+                            self.logger.info(f"  Found embedded image block (Page {page_num+1}, Block {block_num+1})")
+                            img_data = block['image']
+
+                            try:
+                                image_bytes= io.BytesIO(img_data)
+                                pil_image = Image.open(image_bytes)
+                                ocr_text_from_block_image = self.ocr_page_with_nanonets_s(
+                                    pil_image, image_bytes, max_new_tokens=15000
+                                )
 
-        Returns:
-            list: A list of strings, where each string is the comprehensive text
-                  for a corresponding page. Returns an empty list if an error occurs.
-        """
+                                if ocr_text_from_block_image.strip():
+                                    self.logger.info("    OCR found text in embedded image block.")
+                                    current_page_markdown_blocks.append(f"\n\n\n{ocr_text_from_block_image.strip()}\n\n")
+                                else:
+                                    self.logger.info(f"    No OCR text from embedded image block. Adding generic placeholder.")
+                                    current_page_markdown_blocks.append("\n\n![Image Placeholder](image_on_page_{page_num+1}_block_{block_num+1}.png)\n\n") # Consider saving images
+                            except Exception as e:
+                                self.logger.error(f"    Error processing embedded image block for OCR: {e}")
+                                current_page_markdown_blocks.append("\n\n![Image Processing Error](error_on_page_{page_num+1}_block_{block_num+1}.png)\n\n")
+
+
+                    # Insert tables at their approximate positions (after blocks are processed for the page)
+                    # You might need more sophisticated logic here if table positions are granular
+                    while (
+                        table_index < len(tables)
+                        and tables[table_index]["page"] == page.number
+                    ):
+                        current_page_markdown_blocks.append(
+                            self.table_to_markdown(tables[table_index]["content"])
+                        )
+                        table_index += 1
 
-        """taken from self:
-            pdf_path (str): The path to the input PDF file.
-            output_path (str): Directory to save debug output (like rendered images).
-        """
+                    # Phase 2: Full-page OCR if the page seems to be a scanned image or lacks sufficient searchable text
+                    # We prioritize actual searchable text and embedded image OCR.
+                    # Only if very little or no text was found, we resort to full-page OCR.
+                    combined_current_page_text_length = len("".join(current_page_markdown_blocks).strip())
 
-        all_pages_text = []
-        the_text= ""
+                    # A heuristic: if almost no searchable text and no significant OCR from embedded images
+                    if not page_has_searchable_text and combined_current_page_text_length < 100: # Threshold for considering "minimal text"
+                        self.logger.info(f"  Page {page_num + 1} appears to be a scanned image or has minimal text. Attempting full-page OCR.")
+                        try:
+                            pix = page.get_pixmap(matrix=fitz.Matrix(300/72, 300/72))
+                            img_bytes = pix.tobytes("png")
+                            image_bytestream= io.BytesIO(img_bytes)
+                            pil_image = Image.open(image_bytestream)
 
-        try:
-            doc = fitz.open(self.pdf_path)
-            logging.info(f"Opened PDF: {self.pdf_path}")
-
-            tables = self.extract_tables()
-            table_index = 0
-            list_counter = 0
-            in_code_block = False
-            code_block_content = ""
-            code_block_lang = None
-            prev_line = ""
-
-            for page_num, page in enumerate(doc):
-                page_text_content = []
-                page_has_searchable_text = False
-
-                logging.info(f"\nProcessing page {page_num + 1}...")
-
-                # --- Phase 1: Extract text from direct text blocks and process embedded images ---
-                blocks = page.get_text('dict')['blocks']
-                text_blocks_content = []
-                image_block_text_content = []
-
-                page_height = page.rect.height
-                links = self.extract_links(page)
-
-                for block_num, block in enumerate(blocks):
-                    if block['type'] == 0:  # Text block
-                        page_has_searchable_text = True
-                        text_blocks_content.append(self.process_text_block(
-                            block,
-                            page_height,
-                            links,
-                            list_counter,
-                            in_code_block,
-                            code_block_content,
-                            code_block_lang,
-                            prev_line,
-                        ))
-
-                        # for line in block['lines']:
-                        #     for span in line['spans']:
-                        #         text_blocks_content.append(span['text'])
-                    elif block['type'] == 1:  # Image block
-                        logging.info(f"  Found embedded image block (Page {page_num+1}, Block {block_num+1})")
-                        img_data = block['image']
-                        img_ext = block['ext']
+                            ocr_text_from_page = self.ocr_page_with_nanonets_s(
+                                pil_image, image_bytestream, max_new_tokens=15000
+                            )
 
-                        try:
-                            # Attempt OCR on the embedded image block
-                            pil_image = Image.open(io.BytesIO(img_data))
-                            # ocr_text_from_block_image = pytesseract.image_to_string(pil_image)
-                            ocr_text_from_block_image= self.ocr_page_with_nanonets_s(pil_image, model, processor, max_new_tokens=15000)
-
-                            if ocr_text_from_block_image.strip():
-                                logging.info(f"    OCR found text in embedded image block.")
-                                image_block_text_content.append(ocr_text_from_block_image.strip())
+                            if ocr_text_from_page.strip():
+                                self.logger.info(f"  Successfully extracted text via full-page OCR for page {page_num + 1}.")
+                                # If full-page OCR yields significant content and other methods didn't,
+                                # replace or augment. Here, we'll replace to avoid double-counting if it's primarily scanned.
+                                # You might choose to append if you want to combine (e.g., if there's header text + scanned body)
+                                if combined_current_page_text_length < 50: # If almost nothing was found before, replace
+                                    current_page_markdown_blocks = [ocr_text_from_page.strip()]
+                                else: # Otherwise, augment (append)
+                                    current_page_markdown_blocks.append(f"\n\n\n{ocr_text_from_page.strip()}\n\n")
                             else:
-                                # If no OCR text, use the caption
-                                # caption = self.caption_image(pil_image)
-                                # if caption:
-                                #     logging.info(f"    No OCR text, using caption for embedded image block.")
-                                #     image_block_text_content.append(caption)
-                                # else:
-                                #     logging.info(f"    No OCR text and no caption for embedded image block.")
-
-                                # a) captioning sucks, b) no need
-                                image_block_text_content.append("An Image")
-
-                        # except pytesseract.TesseractNotFoundError:
-                        #     logging.warning("    Tesseract-OCR not found. Skipping OCR for embedded image block.")
-                            # caption = self.process_image_block(page, block)
-                            # if caption: image_block_text_content.append(caption)
-
-                            # image_block_text_content.append("An Image")
+                                self.logger.info(f"  Full-page OCR yielded no text for page {page_num+1}.")
                         except Exception as e:
-                            logging.error(f"    Error processing embedded image block for OCR/caption: {e}")
-                            # caption = self.process_image_block(page, block)
-                            # if caption: image_block_text_content.append(caption)
-                            image_block_text_content.append("An Image")
-
-
-                # Insert tables at their approximate positions
-                while (
-                    table_index < len(tables)
-                    and tables[table_index]["page"] == page.number
-                ):
-                    page_text_content += (
-                        "\n\n"
-                        + self.table_to_markdown(tables[table_index]["content"])
-                        + "\n\n"
-                    )
-                    table_index += 1
-
-                # Add content from text blocks
-                if text_blocks_content:
-                    page_text_content.append(" ".join(text_blocks_content))
-
-                # Add content from image blocks
-                if image_block_text_content:
-                    page_text_content.append("\n".join(image_block_text_content))
-
-
-                # --- Phase 2: OCR the entire page IF it seems to be a scanned image ---
-                # We check if page_has_searchable_text is False or if the amount of text
-                # is very small, suggesting it might be mostly a scanned page.
-                # A threshold of 50 characters is arbitrary; adjust as needed.
-                current_text_len = len(" ".join(page_text_content).strip())
-
-                if not page_has_searchable_text or current_text_len < 50:
-                    logging.info(f"  Page {page_num + 1} appears to be a scanned image or has minimal text. Attempting full-page OCR.")
-                    try:
-                        # Render the page as a high-resolution image (e.g., 300 DPI)
-                        pix = page.get_pixmap(matrix=fitz.Matrix(300/72, 300/72))
-                        img_bytes = pix.tobytes("png")
-
-                        pil_image = Image.open(io.BytesIO(img_bytes))
-
-                        # Perform OCR on the entire page image
-                        # ocr_text_from_page = pytesseract.image_to_string(pil_image)
-                        ocr_text_from_page= self.ocr_page_with_nanonets_s(pil_image, model, processor, max_new_tokens=15000)
-
-                        if ocr_text_from_page.strip():
-                            logging.info(f"  Successfully extracted text via full-page OCR.")
-                            page_text_content.append(ocr_text_from_page.strip())
-                        else:
-                            logging.info(f"  Full-page OCR yielded no text for page {page_num+1}.")
-
-                    # except pytesseract.TesseractNotFoundError:
-                    #     logging.warning("  Tesseract-OCR not found. Skipping full-page OCR for this page.")
-                    except Exception as e:
-                        logging.error(f"  Error during full-page OCR on page {page_num+1}: {e}")
-                else:
-                    logging.info(f"  Page {page_num + 1} has sufficient searchable text; skipping full-page OCR.")
-
-
-                # Concatenate all collected text for the current page
-                final_page_text = "\n".join(filter(None, page_text_content)).strip() # Use filter(None, ...) to remove empty strings
-                all_pages_text.append(self.post_process_markdown(final_page_text))
-                the_text += final_page_text + self.page_delimiter
-
-                logging.info(f"  Comprehensive text for page {page_num + 1} (first 200 chars):\n{final_page_text[:200]}...")
+                            self.logger.error(f"  Error during full-page OCR on page {page_num+1}: {e}")
+                    else:
+                        self.logger.info(f"  Page {page_num + 1} has sufficient searchable text or embedded image OCR; skipping full-page OCR.")
 
-                print("\npage done\n")
-                print(final_page_text)
+                    # Join collected markdown blocks for the current page
+                    final_page_markdown = "\n".join(filter(None, current_page_markdown_blocks)).strip()
+                    all_pages_markdown.append(self.post_process_markdown(final_page_markdown))
+                    full_document_markdown.append(self.post_process_markdown(final_page_markdown))
+                    full_document_markdown.append(self.page_delimiter)
 
 
-            doc.close()
-            return the_text, all_pages_text
+                    self.logger.info(f"  Comprehensive text for page {page_num + 1} (first 200 chars):\n{final_page_markdown[:200]}...")
+                    print(f"\n--- Page {page_num+1} Done ---\n")
+                    print(final_page_markdown[:500]) # Print first 500 chars of page markdown
 
-        except fitz.FileNotFoundError:
-            logging.error(f"PDF file not found: {self.pdf_path}")
-            return []
-        except Exception as e:
-            logging.critical(f"An unexpected error occurred: {e}")
-            return []
+                doc.close()
+                return "".join(full_document_markdown), all_pages_markdown
 
+            except fitz.FileNotFoundError:
+                self.logger.error(f"PDF file not found: {self.pdf_path}")
+                return "", []
+            except Exception as e:
+                self.logger.critical(f"An unexpected error occurred during markdown extraction: {e}")
+                self.logger.exception(traceback.format_exc())
+                return "", []
 
     def extract_tables(self):
         """Extract tables from PDF using pdfplumber."""
@@ -449,13 +374,13 @@ class MarkdownPDFExtractor(PDFExtractor):
             self.logger.exception(traceback.format_exc())
             return ""
 
-    def perform_ocr(self, image):
+    def perform_ocr(self, image, image_bytes):
         """Perform OCR on the given image."""
         try:
             # ocr_result = pytesseract.image_to_string(
             #     image
             # )
-            ocr_result= self.ocr_page_with_nanonets_s(image, model, processor, max_new_tokens=15000)
+            ocr_result= self.ocr_page_with_nanonets_s(image, image_bytes, max_new_tokens=15000)
 
 
             return ocr_result.strip()
@@ -464,10 +389,10 @@ class MarkdownPDFExtractor(PDFExtractor):
             self.logger.exception(traceback.format_exc())
             return ""
 
-    def caption_image(self, image):
+    def caption_image(self, image, image_bytes):
         """Generate a caption for the given image."""
         try:
-            ocr_text = self.perform_ocr(image)
+            ocr_text = self.perform_ocr(image, image_bytes)
             if ocr_text:
                 return ocr_text
 
@@ -475,19 +400,38 @@ class MarkdownPDFExtractor(PDFExtractor):
             if image.mode != "RGB":
                 image = image.convert("RGB")
 
-            # Ensure the image is in the correct shape
-            image = np.array(image).transpose(2, 0, 1)  # Convert to (C, H, W) format
+            image_format = image.format
+            dummy_filename = f"dummy.{image_format.lower()}"
+            mime_type, _ = mimetypes.guess_type(dummy_filename)
+
+            if "gemini" in self.MODEL_NAME:
+                response=  self.gclient.models.generate_content(
+                    model= self.MODEL_NAME,
+                    contents=[
+                        types.Part.from_bytes(
+                            data=image_bytes.getvalue(),
+                            mime_type= mime_type
+                        ),
+                        "Write a caption for this image"
+                    ]
+                )
+                return response.text
+            else:
+                # Ensure the image is in the correct shape
+                image = np.array(image).transpose(2, 0, 1)  # Convert to (C, H, W) format
 
-            inputs = self.feature_extractor(images=image, return_tensors="pt").to(
-                self.device
-            )
-            pixel_values = inputs.pixel_values
+                inputs = self.feature_extractor(images=image, return_tensors="pt").to(
+                    self.device
+                )
+                pixel_values = inputs.pixel_values
 
-            generated_ids = self.model.generate(pixel_values, max_length=30)
-            generated_caption = self.tokenizer.batch_decode(
-                generated_ids, skip_special_tokens=True
-            )[0]
-            return generated_caption.strip()
+                generated_ids = self.model.generate(pixel_values, max_length=30)
+
+                generated_ids = self.model.generate(pixel_values, max_length=30)
+                generated_caption = self.tokenizer.batch_decode(
+                    generated_ids, skip_special_tokens=True
+                )[0]
+                return generated_caption.strip()
         except Exception as e:
             self.logger.error(f"Error captioning image: {e}")
             self.logger.exception(traceback.format_exc())
@@ -789,7 +733,11 @@ class MarkdownPDFExtractor(PDFExtractor):
                 Path(self.output_path) / image_filename
             )  # Convert to Path object
             image.save(image_path, "PNG", optimize=True, quality=95)
-            caption = self.caption_image(image)
+
+            img_byte_arr = io.BytesIO()
+            image.save(img_byte_arr)
+            caption = self.caption_image(image, img_byte_arr)
+
             if not caption:
                 caption = (
                     f"{self.pdf_filename}_image_{int(page.number)+1}_{block['number']}"

{pembot-0.0.3.dist-info → pembot-0.0.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pembot
-Version: 0.0.3
+Version: 0.0.4
 Summary: A Python Package to convert PEM blog content to usseful information by leveraging LLMs
 Author-email: cyto <aryan_sidhwani@protonmail.com>
 License-Expression: MIT

{pembot-0.0.3.dist-info → pembot-0.0.4.dist-info}/RECORD

@@ -1,8 +1,8 @@
 pembot/.gitignore,sha256=_7FTsZokJ_pzEyyPjOsGw5x5Xx3gUBFaafs7UlPsv9E,98
 pembot/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-pembot/__init__.py,sha256=2Q6n9cB2MiCO_Fysy8FImzSyQPY4ZWcHfoUGrWP8y8M,211
+pembot/__init__.py,sha256=gDyRo0zcHlKqOSZfoLopwIZ7zXGBXlogi4tihmIeC5o,211
 pembot/gartner.py,sha256=3ALknQ5mSXIimmwCa3JFDzB_EW2hHEcQO1T2odyBquk,5408
-pembot/main.py,sha256=qjfqM_vHRZsnvU7DtvR2dJNUy1B6oOyGJ0LAe1sIOu4,9052
+pembot/main.py,sha256=lZLIV8XPonvNoY4LVS-5fct1y9URMXWoSGJUKMw3Yg8,9667
 pembot/output_structure_local.py,sha256=YfpHzfTNeLMSsB_CjAamha9D6Iz7E1IC-tW9xPCMWFc,3000
 pembot/pem.py,sha256=mv6iGcN1peSY7z2dtCQ_BKj31EFBNfczBhps_d-0XDo,6377
 pembot/query.py,sha256=EDbvQ4It1S8Zg4y-sUCLZjiyBmvQGUAMOVq9yPs_hdE,8242
@@ -79,17 +79,17 @@ pembot/.git/refs/heads/main,sha256=mQvXNyC6knHS2Siuig65zcGaLxzAR8zX_Zoqw2Oj6S0,4
 pembot/.git/refs/remotes/origin/HEAD,sha256=K7aiSqD8bEhBAPXVGim7rYQc0sdV9dk_qiBOXbtOsrQ,30
 pembot/.git/refs/remotes/origin/main,sha256=mQvXNyC6knHS2Siuig65zcGaLxzAR8zX_Zoqw2Oj6S0,41
 pembot/AnyToText/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pembot/AnyToText/convertor.py,sha256=0_ogFWvZEl2x3t3W8-sHjmUUCJe2-Nuxjp_SjTb-mqc,13098
+pembot/AnyToText/convertor.py,sha256=26Pq4OLhVNHgIhJdLLcxGPFTtdnG2lsQkR_53_zkZZM,16997
 pembot/TextEmbedder/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pembot/TextEmbedder/gemini_embedder.py,sha256=P679-2mmQESlYKML1vcrwx_-CSgWJgIQk7NL4F7BLQE,677
 pembot/TextEmbedder/mongodb_embedder.py,sha256=pD8mP-uC_o0COPdOrCTMpoC5PdF8hXlqARHvTr2T-VI,9642
 pembot/TextEmbedder/mongodb_index_creator.py,sha256=ejpsF_y1zY6Z0nux02vjODiDPnxx-YA_xy2PmT94zZ4,5306
 pembot/TextEmbedder/vector_query.py,sha256=Kh1uhx9CatB-oQlQtnW-1I2Qz7MGHI20n2h_8peAChM,1986
-pembot/config/config.yaml,sha256=tM1q7yXsVKJTUDya1JMaclNkAbpTgmcn_7uWiLU-8V8,156
+pembot/config/config.yaml,sha256=IJG9CGEPxH83Cr4X5OEZj9SLFT4ZfwINrA63GpjQssA,156
 pembot/pdf2markdown/LICENSE,sha256=1JTJhQjUYDqJzFJhNtitm7mHyE71PRHgetIqRRWg6Pk,1068
 pembot/pdf2markdown/README.md,sha256=jitM1pwI69oa0N4mXv5-SY1ka9Sz3jsRNCDdpW-50kY,4545
 pembot/pdf2markdown/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pembot/pdf2markdown/extract.py,sha256=4AtYbpsWfQw7w-oU-ue4IL_jdNC4zpA-OjPHHXC2Jdo,35648
+pembot/pdf2markdown/extract.py,sha256=LV9cqkWWmQk-32hsfPdjzbzVVF0kYxzyH7-ssUALxpU,34956
 pembot/pdf2markdown/requirements.txt,sha256=0vZQzkSZKLNVUttd4euoDyYEy0nc2W3CIVxhepHW5Ho,76
 pembot/pdf2markdown/.git/HEAD,sha256=KNJb-Cr0wOK3L1CVmyvrhZ4-YLljCl6MYD2tTdsrboA,21
 pembot/pdf2markdown/.git/config,sha256=ltEWI476vFz2goGWD7QmCDvC6UCQ9ELviXuURlvte_w,269
@@ -123,7 +123,7 @@ pembot/pdf2markdown/config/config.yaml,sha256=w75W2Eg4-tu8rRk_23PqxWDh0010kRKLmP
 pembot/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pembot/utils/inference_client.py,sha256=jeURmY2P5heVlH1dCV0XSgiX3U2qYGEmrnUv0KFpdww,5380
 pembot/utils/string_tools.py,sha256=gtRa5rBR0Q7GspTu2WtCnvhJQLFjPfWLvhmyiPkyStU,1883
-pembot-0.0.3.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-pembot-0.0.3.dist-info/WHEEL,sha256=Dyt6SBfaasWElUrURkknVFAZDHSTwxg3PaTza7RSbkY,100
-pembot-0.0.3.dist-info/METADATA,sha256=lhYkhIOKuUxiOYnD7Ms28SWjHMJQE6rfLAshsyQ_JBk,313
-pembot-0.0.3.dist-info/RECORD,,
+pembot-0.0.4.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+pembot-0.0.4.dist-info/WHEEL,sha256=Dyt6SBfaasWElUrURkknVFAZDHSTwxg3PaTza7RSbkY,100
+pembot-0.0.4.dist-info/METADATA,sha256=v9TTMdInby4h5K0aKMrT11NgpyBw0ZMpzrC9dgI0O14,313
+pembot-0.0.4.dist-info/RECORD,,