ai-parrot 0.17.2__cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

parrot/loaders/audio.py

@@ -0,0 +1,199 @@
+from typing import Any, List
+from collections.abc import Callable
+from pathlib import PurePath
+from ..stores.models import Document
+from .basevideo import BaseVideoLoader
+
+
+class AudioLoader(BaseVideoLoader):
+    """
+    Generating transcripts from local Audio.
+    """
+    extensions: List[str] = ['.mp3', '.webm', '.ogg']
+
+    def load_video(self, path):
+        return
+
+    async def load_audio(self, path: PurePath) -> list:
+        metadata = {
+            "source": f"{path}",
+            "url": f"{path.name}",
+            # "index": path.stem,
+            "filename": f"{path}",
+            'type': 'audio_transcript',
+            "source_type": self._source_type,
+            "document_meta": {
+                "language": self._language,
+                "topic_tags": ""
+            }
+        }
+        documents = []
+
+        # Paths for outputs
+        vtt_path = path.with_suffix(".vtt")
+        txt_path = path.with_suffix(".txt")
+        srt_path = path.with_suffix(".srt")
+
+        # ensure a clean 16k Hz mono wav file for whisper
+        wav_path = self.ensure_wav_16k_mono(path)
+        # get the Whisper parser
+        transcript_whisper = self.get_whisperx_transcript(wav_path)
+        transcript = transcript_whisper.get('text', '') if transcript_whisper else ''
+        try:
+            self.saving_file(txt_path, transcript.encode("utf-8"))
+            print(f"Saved TXT transcript to: {txt_path}")
+        except Exception as exc:
+            print(f"Error saving TXT transcript: {exc}")
+        if transcript:
+            doc = Document(
+                page_content=transcript,
+                metadata={
+                    "source": f"{txt_path}",
+                    "url": f"{txt_path.name}",
+                    "filename": f"{txt_path}",
+                    "origin": f"{path}",
+                    'type': 'audio_transcript',
+                    "source_type": 'AUDIO',
+                }
+            )
+        # diarization:
+        if self._diarization:
+            if (srt := self.audio_to_srt(
+                audio_path=wav_path,
+                asr=transcript_whisper,
+                output_srt_path=srt_path,
+                max_gap_s=0.5,
+                max_chars=90,
+                max_duration_s=0.9,
+            )):
+                doc = Document(
+                    page_content=srt,
+                    metadata={
+                        "source": f"{srt_path}",
+                        "url": f"{srt_path.name}",
+                        "filename": f"{srt_path}",
+                        "origin": f"{path}",
+                        'type': 'srt_transcript',
+                        "source_type": 'AUDIO',
+                    }
+                )
+        # Summarize the transcript (only if enabled)
+        if self._summarization and transcript:
+            try:
+                summary = await self.summary_from_text(transcript)
+                # Create Two Documents, one is for transcript, second is VTT:
+                doc = Document(
+                    page_content=summary,
+                    metadata={
+                        "source": f"{path}",
+                        "url": f"{path.name}",
+                        "filename": f"{path}",
+                        "origin": f"{path}",
+                        'type': 'summary',
+                        "source_type": 'TEXT',
+                    }
+                )
+                documents.append(doc)
+            except Exception as exc:
+                print(f"Error generating summary: {exc}")
+        if transcript_whisper:
+            # VTT version:
+            vtt_text = self.transcript_to_vtt(transcript_whisper, vtt_path)
+            doc = Document(
+                page_content=vtt_text,
+                metadata={
+                    "source": f"{vtt_path}",
+                    "url": f"{vtt_path.name}",
+                    "filename": f"{vtt_path}",
+                    "origin": f"{path}",
+                    'type': 'vtt_transcript',
+                    "source_type": 'TEXT',
+                }
+            )
+            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
+
+    async def extract_audio(self, path: PurePath) -> list:
+        """
+        Extract audio transcript and summary from a local audio file.
+        """
+        vtt_path = path.with_suffix('.vtt')
+        transcript_path = path.with_suffix('.txt')
+        srt_path = path.with_suffix('.srt')
+        summary_path = path.with_suffix('.summary')
+        metadata = {
+            "source": f"{path}",
+            "url": f"{path.name}",
+            "filename": f"{path}",
+            'type': 'audio_transcript',
+            "source_type": self._source_type,
+            "vtt_path": f"{vtt_path}",
+            "transcript_path": f"{transcript_path}",
+            "srt_path": f"{srt_path}",
+            "summary_path": f"{summary_path}",
+            "document_meta": {
+                "language": self._language,
+            }
+        }
+        # get the Whisper parser
+        # ensure a clean 16k Hz mono wav file for whisper
+        wav_path = self.ensure_wav_16k_mono(path)
+        # get the Whisper parser
+        transcript_whisper = self.get_whisperx_transcript(wav_path)
+        if transcript_whisper:
+            transcript = transcript_whisper['text']
+        else:
+            transcript = ''
+        if self._diarization:
+            srt = self.audio_to_srt(
+                audio_path=wav_path,
+                asr=transcript_whisper,
+                output_srt_path=srt_path,
+                max_gap_s=0.5,
+                max_chars=90,
+                max_duration_s=0.9,
+            )
+            if srt:
+                try:
+                    self.saving_file(srt_path, srt.encode("utf-8"))
+                    print(f"Saved SRT transcript to: {srt_path}")
+                except Exception as exc:
+                    print(f"Error saving SRT transcript: {exc}")
+        # Summarize the transcript (only if enabled)
+        self.saving_file(transcript_path, transcript.encode('utf-8'))
+        if self._summarization and transcript:
+            try:
+                summary = await self.summary_from_text(transcript)
+                # Create Two Documents, one is for transcript, second is VTT:
+                metadata['summary'] = summary
+                self.saving_file(summary_path, summary.encode('utf-8'))
+            except Exception as exc:
+                print(f"Error generating summary: {exc}")
+            # VTT version:
+            transcript = self.transcript_to_vtt(transcript_whisper, vtt_path)
+        return metadata
+
+    async def _load(self, source, **kwargs) -> List[Document]:
+        return await self.load_audio(source)

parrot/loaders/basepdf.py

@@ -0,0 +1,53 @@
+from abc import abstractmethod
+from typing import List, Union, Callable, Optional
+from pathlib import Path
+from pathlib import PurePath
+from .abstract import AbstractLoader
+from ..stores.models import Document
+
+
+class BasePDF(AbstractLoader):
+    """
+    Base Abstract loader for all PDF-file Loaders.
+    """
+    extensions: set[str] = {'.pdf'}
+
+    def __init__(
+        self,
+        source: Optional[Union[str, Path, List[Union[str, Path]]]] = None,
+        *,
+        tokenizer: Union[str, Callable] = None,
+        text_splitter: Union[str, Callable] = None,
+        source_type: str = 'file',
+        as_markdown: bool = False,
+        use_chapters: bool = False,
+        use_pages: bool = False,
+        **kwargs
+    ):
+        super().__init__(
+            source,
+            tokenizer=tokenizer,
+            text_splitter=text_splitter,
+            source_type=source_type,
+            **kwargs
+        )
+        self._lang = 'eng'
+        self.doctype = 'pdf'
+        self._source_type = source_type
+        self.as_markdown = as_markdown
+        self.use_chapters = use_chapters
+        self.use_pages = use_pages
+
+    @abstractmethod
+    async def _load(self, path: Union[str, PurePath, List[PurePath]], **kwargs) -> List[Document]:
+        """Load data from a source and return it as a Langchain Document.
+
+        Args:
+            path (Union[str, PurePath, List[PurePath]]): The source of the data.
+
+        Returns:
+            List[Document]: A list of Langchain Documents.
+        """
+        self.logger.info(
+            f"Loading file: {path}"
+        )