ai-parrot 0.3.4__cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

parrot/loaders/audio.py

@@ -0,0 +1,106 @@
+from typing import Any
+from collections.abc import Callable
+from pathlib import PurePath
+from langchain.docstore.document import Document
+from .basevideo import BaseVideoLoader
+
+
+class AudioLoader(BaseVideoLoader):
+    """
+    Generating transcripts from local Audio.
+    """
+    _extension = ['.mp3', '.webm']
+
+    def __init__(
+        self,
+        path: PurePath,
+        tokenizer: Callable[..., Any] = None,
+        text_splitter: Callable[..., Any] = None,
+        source_type: str = 'documentation',
+        encoding: str = 'utf-8',
+        origin: str = '',
+        **kwargs
+    ):
+        super().__init__(tokenizer, text_splitter, source_type=source_type, **kwargs)
+        self.path = path
+
+    def load_audio(self, path: PurePath) -> list:
+        metadata = {
+            "source": f"{path}",
+            "url": f"{path.name}",
+            "index": path.stem,
+            "filename": f"{path}",
+            "question": '',
+            "answer": '',
+            'type': 'audio_transcript',
+            "source_type": self._source_type,
+            "summary": '',
+            "document_meta": {
+                "language": self._language,
+                "topic_tags": ""
+            }
+        }
+        documents = []
+        transcript_path = path.with_suffix('.vtt')
+        # get the Whisper parser
+        transcript_whisper = self.get_whisper_transcript(path)
+        if transcript_whisper:
+            transcript = transcript_whisper['text']
+        else:
+            transcript = ''
+        # Summarize the transcript
+        if transcript:
+            summary = self.get_summary_from_text(transcript)
+            # Create Two Documents, one is for transcript, second is VTT:
+            metadata['summary'] = summary
+            doc = Document(
+                page_content=transcript,
+                metadata=metadata
+            )
+            documents.append(doc)
+        if transcript_whisper:
+            # VTT version:
+            transcript = self.transcript_to_vtt(transcript_whisper, transcript_path)
+            doc = Document(
+                page_content=transcript,
+                metadata=metadata
+            )
+            documents.append(doc)
+            # Saving every dialog chunk as a separate document
+            dialogs = self.transcript_to_blocks(transcript_whisper)
+            docs = []
+            for chunk in dialogs:
+                _meta = {
+                    "index": f"{path.stem}:{chunk['id']}",
+                    "document_meta": {
+                        "start": f"{chunk['start_time']}",
+                        "end": f"{chunk['end_time']}",
+                        "id": f"{chunk['id']}",
+                        "language": self._language,
+                        "title": f"{path.stem}",
+                        "topic_tags": ""
+                    }
+                }
+                _info = {**metadata, **_meta}
+                doc = Document(
+                    page_content=chunk['text'],
+                    metadata=_info
+                )
+                docs.append(doc)
+            documents.extend(docs)
+        return documents
+
+    def load(self) -> list:
+        documents = []
+        if self.path.is_file():
+            docs = self.load_audio(self.path)
+            documents.extend(docs)
+        if self.path.is_dir():
+            # iterate over the files in the directory
+            for ext in self._extension:
+                for item in self.path.glob(f'*{ext}'):
+                    if set(item.parts).isdisjoint(self.skip_directories):
+                        documents.extend(
+                            self.load_audio(item)
+                        )
+        return self.split_documents(documents)

parrot/loaders/basepdf.py

@@ -0,0 +1,102 @@
+from collections.abc import Callable
+from typing import Any
+from abc import abstractmethod
+from pathlib import Path, PurePath
+from PIL import Image
+from .abstract import AbstractLoader
+from ..conf import STATIC_DIR
+
+
+class BasePDF(AbstractLoader):
+    """
+    Base Abstract loader for all PDF files.
+    """
+    _extension = ['.pdf']
+
+    def __init__(
+        self,
+        path: PurePath,
+        tokenizer: Callable[..., Any] = None,
+        text_splitter: Callable[..., Any] = None,
+        source_type: str = 'pdf',
+        language: str = "eng",
+        **kwargs
+    ):
+        super().__init__(tokenizer, text_splitter, source_type=source_type, **kwargs)
+        self.path = path
+        if isinstance(path, str):
+            self.path = Path(path).resolve()
+        self.save_images: bool = bool(kwargs.get('save_images', False))
+        self._imgdir = STATIC_DIR.joinpath('images')
+        if self.save_images is True:
+            if self._imgdir.exists() is False:
+                self._imgdir.mkdir(parents=True, exist_ok=True)
+        if language == 'en':
+            language = 'eng'
+        self._lang = language
+
+    def save_image(self, img_stream: Image, image_name: str, save_path: Path):
+        # Create the image directory if it does not exist
+        if save_path.exists() is False:
+            save_path.mkdir(parents=True, exist_ok=True)
+        img_path = save_path.joinpath(image_name)
+        self.logger.notice(
+            f"Saving Image Page on {img_path}"
+        )
+        if not img_path.exists():
+            # Save the image
+            img_stream.save(img_path, format="PNG", optimize=True)
+        return img_path
+
+    @abstractmethod
+    def _load_pdf(self, path: Path) -> list:
+        """
+        Load a PDF file using Fitz.
+
+        Args:
+            path (Path): The path to the PDF file.
+
+        Returns:
+            list: A list of Langchain Documents.
+        """
+        pass
+
+    def load(self) -> list:
+        """
+        Load data from a PDF file.
+
+        Args:
+            source (str): The path to the PDF file.
+
+        Returns:
+            list: A list of Langchain Documents.
+        """
+        if isinstance(self.path, list):
+            # list of files:
+            documents = []
+            for p in self.path:
+                if self._check_path(p):
+                    documents.extend(self._load_pdf(p))
+        if not self.path.exists():
+            raise FileNotFoundError(
+                f"PDF file/directory not found: {self.path}"
+            )
+        if self.path.is_dir():
+            documents = []
+            # iterate over the files in the directory
+            for ext in self._extension:
+                for item in self.path.glob(f'*{ext}'):
+                    if set(item.parts).isdisjoint(self.skip_directories):
+                        documents.extend(self._load_pdf(item))
+        elif self.path.is_file():
+            documents = self._load_pdf(self.path)
+        else:
+            raise ValueError(
+                f"PDF Loader: Invalid path: {self.path}"
+            )
+        return self.split_documents(documents)
+
+    def parse(self, source):
+        raise NotImplementedError(
+            "Parser method is not implemented for PDFLoader."
+        )

parrot/loaders/basevideo.py

@@ -0,0 +1,280 @@
+from collections.abc import Callable
+from typing import Any, Union, List
+from abc import abstractmethod
+from pathlib import Path
+from moviepy.editor import VideoFileClip
+from transformers import (
+    pipeline,
+    AutoModelForSeq2SeqLM,
+    AutoTokenizer,
+    GenerationConfig
+)
+from transformers import (
+    WhisperForConditionalGeneration,
+    WhisperProcessor,
+    WhisperTimeStampLogitsProcessor
+)
+
+from langchain.chains.summarize import load_summarize_chain
+from langchain.text_splitter import (
+    TokenTextSplitter,
+)
+from .abstract import AbstractLoader
+
+
+def extract_video_id(url):
+    parts = url.split("?v=")
+    video_id = parts[1].split("&")[0]
+    return video_id
+
+
+class BaseVideoLoader(AbstractLoader):
+    """
+    Generating Video transcripts from Videos.
+    """
+    _extension = ['.youtube']
+    encoding = 'utf-8'
+    chunk_size = 768
+
+    def __init__(
+        self,
+        urls: List[str],
+        tokenizer: Callable[..., Any] = None,
+        text_splitter: Callable[..., Any] = None,
+        source_type: str = 'video',
+        language: str = "en",
+        video_path: Union[str, Path] = None,
+        **kwargs
+    ):
+        super().__init__(tokenizer, text_splitter, source_type, **kwargs)
+        self.urls = urls
+        self._task = kwargs.get('task', "automatic-speech-recognition")
+        # Topics:
+        self.topics: list = kwargs.get('topics', [])
+        self._model_size: str = kwargs.get('model_size', 'medium')
+        self.summarization_model = "facebook/bart-large-cnn"
+        self._model_name: str = kwargs.get('model_name', 'whisper')
+        self.summarizer = pipeline(
+            "summarization",
+            tokenizer=AutoTokenizer.from_pretrained(
+                self.summarization_model
+            ),
+            model=AutoModelForSeq2SeqLM.from_pretrained(
+                self.summarization_model
+            ),
+            device=self._get_device()
+        )
+        # language:
+        self._language = language
+        # directory:
+        if isinstance(video_path, str):
+            self._video_path = Path(video_path).resolve()
+        self._video_path = video_path
+
+    def transcript_to_vtt(self, transcript: str, transcript_path: Path) -> str:
+        """
+        Convert a transcript to VTT format.
+        """
+        vtt = "WEBVTT\n\n"
+        for i, chunk in enumerate(transcript['chunks'], start=1):
+            start, end = chunk['timestamp']
+            text = chunk['text'].replace("\n", " ")  # Replace newlines in text with spaces
+
+            if start is None or end is None:
+                print(f"Warning: Missing timestamp for chunk {i}, skipping this chunk.")
+                continue
+
+            # Convert timestamps to WebVTT format (HH:MM:SS.MMM)
+            start_vtt = f"{int(start // 3600):02}:{int(start % 3600 // 60):02}:{int(start % 60):02}.{int(start * 1000 % 1000):03}"
+            end_vtt = f"{int(end // 3600):02}:{int(end % 3600 // 60):02}:{int(end % 60):02}.{int(end * 1000 % 1000):03}"
+
+            vtt += f"{i}\n{start_vtt} --> {end_vtt}\n{text}\n\n"
+        # Save the VTT file
+        try:
+            with open(str(transcript_path), "w") as f:
+                f.write(vtt)
+            print(f'Saved VTT File on {transcript_path}')
+        except Exception as exc:
+            print(f"Error saving VTT file: {exc}")
+        return vtt
+
+    def format_timestamp(self, seconds):
+        # This helper function takes the total seconds and formats it into hh:mm:ss,ms
+        hours, remainder = divmod(int(seconds), 3600)
+        minutes, seconds = divmod(remainder, 60)
+        milliseconds = int((seconds % 1) * 1000)
+        seconds = int(seconds)
+        return f"{hours:02}:{minutes:02}:{seconds:02},{milliseconds:03}"
+
+    def transcript_to_blocks(self, transcript: str) -> list:
+        """
+        Convert a transcript to blocks.
+        """
+        blocks = []
+        for i, chunk in enumerate(transcript['chunks'], start=1):
+            current_window = {}
+            start, end = chunk['timestamp']
+            if start is None or end is None:
+                print(f"Warning: Missing timestamp for chunk {i}, skipping this chunk.")
+                continue
+
+            start_srt = self.format_timestamp(start)
+            end_srt = self.format_timestamp(end)
+            text = chunk['text'].replace("\n", " ")  # Replace newlines in text with spaces
+            current_window['id'] = i
+            current_window['start_time'] = start_srt
+            current_window['end_time'] = end_srt
+            current_window['text'] = text
+            blocks.append(current_window)
+        return blocks
+
+    def transcript_to_srt(self, transcript: str) -> str:
+        """
+        Convert a transcript to SRT format.
+        """
+        # lines = transcript.split("\n")
+        srt = ""
+        for i, chunk in enumerate(transcript['chunks'], start=1):
+            start, end = chunk['timestamp']
+            text = chunk['text'].replace("\n", " ")  # Replace newlines in text with spaces
+            # Convert start and end times to SRT format HH:MM:SS,MS
+            start_srt = f"{start // 3600:02}:{start % 3600 // 60:02}:{start % 60:02},{int(start * 1000 % 1000):03}"
+            end_srt = f"{end // 3600:02}:{end % 3600 // 60:02}:{end % 60:02},{int(end * 1000 % 1000):03}"
+            srt += f"{i}\n{start_srt} --> {end_srt}\n{text}\n\n"
+        return srt
+
+    def chunk_text(self, text, chunk_size, tokenizer):
+        # Tokenize the text and get the number of tokens
+        tokens = tokenizer.tokenize(text)
+        # Split the tokens into chunks
+        for i in range(0, len(tokens), chunk_size):
+            yield tokenizer.convert_tokens_to_string(
+                tokens[i:i+chunk_size]
+            )
+
+    def get_summary(self, documents: list) -> str:
+        """
+        Get a summary of a text.
+        """
+        try:
+            splitter = TokenTextSplitter(
+                chunk_size=5000,
+                chunk_overlap=100,
+            )
+            summarize_chain = load_summarize_chain(
+                llm=self._llm,
+                chain_type="refine"
+            )
+            chunks = splitter.split_documents(documents)
+            summary = summarize_chain.invoke(chunks)
+            return summary
+        except Exception as e:
+            print('ERROR in get_summary:', e)
+            return ""
+
+    def summarization(self, text: str) -> str:
+        """
+        Get a summary of a text considering token limits.
+        """
+        try:
+            tokenizer = self.summarizer.tokenizer
+            # to be safe under the limit
+            max_length = tokenizer.model_max_length - 10
+            summaries = []
+            for text_chunk in self.chunk_text(text, max_length, tokenizer):
+                chunk_summary = self.summarizer(
+                    text_chunk,
+                    max_length=150,
+                    min_length=30,
+                    do_sample=False)[0]['summary_text']
+                summaries.append(chunk_summary)
+            return " ".join(summaries)
+        except Exception as e:
+            print('ERROR in summarization:', e)
+            return ""
+
+    def extract_audio(self, video_path, audio_path):
+        """
+        Extracts the audio from a video file and saves it as an audio file.
+
+        Args:
+            video_path (str): Path to the video file.
+            audio_path (str): Path where the extracted audio file will be saved.
+        """
+        if audio_path.exists():
+            print(f"Audio already extracted: {audio_path}")
+            return
+        video_clip = VideoFileClip(str(video_path))
+        audio_clip = video_clip.audio
+        if not audio_clip:
+            return
+        audio_clip.write_audiofile(str(audio_path))
+        audio_clip.close()
+        video_clip.close()
+
+    def get_whisper_transcript(self, audio_path: Path, chunk_length: int = 30):
+        # Initialize the Whisper parser
+        if self._model_name == 'whisper':
+            if self._language == 'en':
+                model_name = f"openai/whisper-{self._model_size}.en"
+            elif self._language == 'es':
+                model_name = f"juancopi81/whisper-{self._model_size}-es"
+            else:
+                model_name = "openai/whisper-large-v3"
+        else:
+            model_name = self._model_name
+
+        # Load the model and processor
+        model = WhisperForConditionalGeneration.from_pretrained(model_name)
+        processor = WhisperProcessor.from_pretrained(model_name)
+
+        # Try to load the generation config, fallback to default if it doesn't exist
+        try:
+            generation_config = GenerationConfig.from_pretrained(model_name)
+        except EnvironmentError:
+            print(
+                f"Warning: No generation_config.json found for model {model_name}. Using default configuration."
+            )
+            generation_config = model.generation_config
+
+        # Check and set the no_timestamps_token_id if it doesn't exist
+        if not hasattr(model.config, 'no_timestamps_token_id'):
+            model.config.no_timestamps_token_id = processor.tokenizer.convert_tokens_to_ids('<|notimestamps|>')
+
+        # Define the generation configuration with WhisperTimeStampLogitsProcessor
+        try:
+            model.config.logits_processor = [
+                WhisperTimeStampLogitsProcessor(generation_config)
+            ]
+        except Exception:
+            model.config.logits_processor = [
+                WhisperTimeStampLogitsProcessor(model.config)
+            ]
+
+        whisper_pipe = pipeline(
+            task=self._task,
+            model=model,
+            tokenizer=processor.tokenizer,
+            feature_extractor=processor.feature_extractor,
+            device=self._get_device(),
+            chunk_length_s=chunk_length,
+            use_fast=True
+        )
+        if audio_path.exists() and audio_path.stat().st_size > 0:
+            # Use the parser to extract transcript
+            return whisper_pipe(
+                str(audio_path),
+                return_timestamps=True
+            )
+        return None
+
+    @abstractmethod
+    def load(self) -> list:
+        pass
+
+    @abstractmethod
+    def load_video(self, url: str, video_title: str, transcript: str) -> list:
+        pass
+
+    def parse(self, source):
+        pass

parrot/loaders/csv.py

@@ -0,0 +1,42 @@
+from pathlib import PurePath
+from langchain_community.document_loaders.csv_loader import CSVLoader as CSVL
+from .abstract import AbstractLoader
+
+
+class CSVLoader(AbstractLoader):
+    """
+    Loader for CSV files.
+    """
+    _extension = ['.csv']
+    csv_args: dict = {
+        "delimiter": ",",
+        "quotechar": '"',
+        "escapechar": "\\",
+        "skipinitialspace": False,
+        "lineterminator": "\n",
+        "quoting": 0,
+        "skiprows": 0,
+        "encoding": None
+    }
+
+    def load(self, path: PurePath) -> list:
+        """
+        Load data from a CSV file.
+
+        Args:
+            source (str): The path to the CSV file.
+
+        Returns:
+            list: A list of Langchain Documents.
+        """
+        if self._check_path(path):
+            self.logger.info(f"Loading CSV file: {path}")
+            loader = CSVL(
+                file_path=path,
+                csv_args=self.csv_args,
+                autodetect_encoding=True
+            )
+            documents = loader.load()
+            return self.split_documents(documents)
+        else:
+            return []

parrot/loaders/dir.py

@@ -0,0 +1,37 @@
+import os
+import glob
+from pathlib import PurePath
+from langchain_community.document_loaders import (
+    DirectoryLoader
+)
+
+
+def load_directory(
+    path: PurePath,
+    text_splitter,
+    source_type,
+    file_pattern: str = "**/*.txt"
+):
+    """
+    Load all Text documents from a directory.
+
+    Args:
+        path (str): The path to the directory.
+        text_splitter (TextSplitter): A text splitter object.
+        source_type (str): The type of source.
+
+    Returns:
+        list: A list of documents.
+    """
+    documents = []
+    loader = DirectoryLoader(
+        path=str(path),
+        glob=file_pattern,
+        recursive=True,
+        show_progress=True,
+        use_multithreading=True
+    )
+    documents = loader.load()
+    for doc in documents:
+        doc.metadata['source_type'] = source_type
+    return text_splitter.split_documents(documents)