utilis.py

import os
import json
from typing import List
import assemblyai as aai
from youtube_transcript_api import YouTubeTranscriptApi
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_core.documents.base import Document
from langchain_community.vectorstores import DeepLake
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.prompts.few_shot import FewShotPromptTemplate
from langchain.prompts.prompt import PromptTemplate
from typing import Dict
import uuid
from pdf2image import convert_from_bytes
import pytesseract
from pytesseract import Output


class Processing:
    def __init__(self, dataset_path: str, embedding_model_name: str, chunk_size=500, chunk_overlap=5):
        """
        Parameters:
            dataset_path (str): Path to the dataset in the Vector-DB
            embedding_model_name (str): Name of the HuggingFace model to be used for embeddings
            chunk_size (int): Size of each chunk to be processed
            chunk_overlap (int): Overlap between each chunk

        Initialize embedding model, text splitter, transcriber and Vector-DB
        """
        self.dataset_path = dataset_path
        self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
        self.transcriber = aai.Transcriber()

        self.embedding_model = HuggingFaceEmbeddings(
            model_name=embedding_model_name,
            encode_kwargs={'normalize_embeddings': False}
        )

        self.db = DeepLake(dataset_path=f"hub://{self.dataset_path}",
                           embedding=self.embedding_model,
                           exec_option="compute_engine"
                           )

    def _add_metadata(self, documents: List[Document], url: str, id: str, source: str, file_type: str,
                      course_tag="") -> (List[Document], Dict[str, str]):
        """
        Parameters:
            documents (List[Document]): List of documents to add metadata to
            id (str): ID of the documents
            source (str): Source of the documents
            file_type (str): Type of the documents
            course_tag (str): Tag to identify the course the documents belongs to

        Returns:
            documents (List[Document], Dict[str, str): List of documents with metadata added along with the metadata

        Add metadata to the documents
        """
        metadata = {
            "id": id,
            "source": source,
            "url": url,
            "file_type": file_type,
            "course_tag": course_tag
        }
        for doc in documents:
            doc.metadata = metadata
        return documents, metadata

    def load_pdf(self, text) -> (List[Document], Dict[str, str]):
        """
        Returns:
            documents (List[Document], Dict[str, str): List of documents with metadata added along with the metadata

        Load PDF file, split into chunks and add metadata
        """
        pdf_chunk = self.text_splitter.create_documents([text])
        print("Created document chunks")
        return self._add_metadata(pdf_chunk, url="NaN", id=str(uuid.uuid4()), source="document", file_type="pdf")

    def load_scanned_pdf(self, file) -> str:
        """
        Parameters:
            file (File): Scanned PDF file to be processed

        Returns:
            str: Text extracted from the scanned PDF file

        Extract text from scanned PDF file
        """
        images = convert_from_bytes(file)

        all_text = ""
        for image in images:
            # Perform OCR on the image
            text = pytesseract.image_to_data(image, lang='eng', output_type=Output.DICT)

            # Extract text from the dictionary
            page_text = " ".join(text['text'])
            all_text += page_text + "\n"

        return all_text

    def load_transcript(self, url) -> (List[Document], Dict[str, str]):
        """
        Returns:
            documents (List[Document], Dict[str, str): List of documents with metadata added along with the metadata

        Load transcript, split into chunks and add metadata
        """
        transcript = self.transcriber.transcribe(url)
        print("Transcribed")
        transcript_chunk = self.text_splitter.create_documents([transcript.text])
        print("Created transcript chunks")
        return self._add_metadata(transcript_chunk, url="NaN", id=str(uuid.uuid4()), source="custom_video",
                                  file_type="transcript")

    def load_yt_transcript(self, url) -> (List[Document], Dict[str, str]):
        """
        Returns:
            documents (List[Document], Dict[str, str): List of documents with metadata added along with the metadata

        Load YouTube transcript, split into chunks and add metadata
        """
        if url.startswith("https://www.youtube.com/watch?v="):
            video_id = url.replace("https://www.youtube.com/watch?v=", "")
        else:
            video_id = url.replace("https://youtu.be/", "")

        transcript = YouTubeTranscriptApi.get_transcript(video_id)
        print("Downloaded transcript")
        transcript = [line['text'] for line in transcript]
        transcript_text = ' '.join(transcript)
        yt_transcript_chunk = self.text_splitter.create_documents([transcript_text])
        print("Created YouTube transcript chunks")
        return self._add_metadata(yt_transcript_chunk, url=url, id=video_id, source="youtube", file_type="transcript")

    def upload_to_db(self, documents: List[Document]):
        """
        Parameters:
            documents (List[Document]): List of documents to upload to Vector-DB

        Upload documents to Vector-DB
        """
        print("Embedding and uploading to Vector-DB...")
        self.db.add_documents(documents)
        print("Uploaded to Vector-DB")


class PromptCreate:
    def __init__(self, example_path, save_path):
        self.examples = []
        self.example_prompt = None
        self.few_shot_prompt = None
        self.example_path = example_path
        self.file_name = "example_{i}.json"
        self.save_path = save_path

    def load_examples(self):
        for i in range(1, len(os.listdir(self.example_path)) + 1):
            filename = os.path.join(self.example_path, self.file_name.format(i=i))
            try:
                with open(filename, "r") as json_file:
                    self.examples.append(json.load(json_file))
            except FileNotFoundError:
                print(f"File {filename} not found.")
            except json.JSONDecodeError:
                print(f"Error decoding JSON from file {filename}.")

    def create_prompt_template(self, input_variables, template_string):
        self.example_prompt = PromptTemplate(input_variables=input_variables, template=template_string)

    def create_few_shot_prompt(self, prefix, suffix):
        self.few_shot_prompt = FewShotPromptTemplate(
            examples=self.examples, example_prompt=self.example_prompt, prefix=prefix, suffix=suffix
        )

    def save_prompt(self):
        self.few_shot_prompt.save(self.save_path)