Upload 5 files

main.py

@@ -0,0 +1 @@
+

model.py

@@ -0,0 +1,128 @@
+from langchain.llms import HuggingFacePipeline
+from langchain.embeddings import HuggingFaceInstructEmbeddings
+from langchain.chains import RetrievalQA
+from transformers import (
+    AutoTokenizer,
+    AutoModelForSeq2SeqLM,
+    pipeline,
+    GenerationConfig
+)
+
+class lamini:
+    def __init__(self):
+        pass
+    
+    def load_model(self, task="text2text-generation", **kwargs) -> HuggingFacePipeline:
+        """Returns a pipeline for the model
+        - model: MBZUAI/LaMini-Flan-T5-248M
+        
+        Returns:
+            _type_: _description_
+        """
+        model_id = "MBZUAI/LaMini-Flan-T5-248M"
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
+        gen_config = GenerationConfig.from_pretrained(model_id)
+        
+        max_length = kwargs.get("max_length", 512)
+        temperature = kwargs.get("temperature", 0)
+        top_p = kwargs.get("top_p", 0.95)
+        repetition_penalty = kwargs.get("repetition_penalty", 1.15)
+        
+        pipe = pipeline( 
+                "text2text-generation",
+                model=model,
+                tokenizer=tokenizer,
+                generation_config=gen_config,
+                max_length=max_length,
+                top_p=top_p,
+                temperature=temperature,
+                repetition_penalty=repetition_penalty,
+        )
+        
+        llm = HuggingFacePipeline(pipeline=pipe)
+        return llm
+
+class templates:
+    def __init__(self, llm: HuggingFacePipeline):
+        self.llm = llm
+    
+    def summarize(self, text, **kwargs):
+        """Summarize text
+        
+        Args:
+            text (str): text to summarize
+        
+        Returns:
+            str: summarized text
+        """
+        
+        instruction = "summarize for better understanding: "
+        text = instruction + text
+        return self.llm(text, **kwargs)
+    
+    def generate_tile(self, text, **kwargs):
+        """Generate a title for text
+        
+        Args:
+            text (str): text to generate title for
+        
+        Returns:
+            str: title
+        """
+        
+        instruction = "generate a title for this text: "
+        text = instruction + text
+        return self.llm(text, **kwargs)
+
+class qa_template:
+    def __init__(self, llm):
+        from langchain.chains.retrieval_qa.base import BaseRetrievalQA
+        self.llm = llm
+        self.qa_inf: BaseRetrievalQA
+    
+    def load(self, knowledge_base):
+        """Load knowledge base
+        
+        Args:
+            knowledge_base (str): knowledge base to load
+        
+        Returns:
+            BaseRetrievalQA: (optional to use) returns QA interface
+        """
+        from utils import LangChainChunker
+        from langchain.vectorstores import Chroma
+        from langchain.chains import RetrievalQA
+        
+        embeds = HuggingFaceInstructEmbeddings(model_name="hkunlp/instructor-large")
+        chunker = LangChainChunker(knowledge_base)
+        chunks = chunker.chunker(size=512)
+        db = Chroma.from_texts(chunks, embeds)
+        retriever = db.as_retriever()
+        
+        qa_inf = RetrievalQA.from_chain_type(
+            llm=self.llm, chain_type="stuff", retriever=retriever
+        )
+        
+        self.qa_inf = qa_inf
+        return qa_inf
+    
+    def start_gradio(self, title: str):
+        """Start gradio interface
+        
+        Returns:
+            _type_: _description_
+        """
+        import gradio as gr
+        
+        def interface(msg, history):
+            res = self.qa_inf.run(msg)
+            return str(res)
+        
+        ui = gr.ChatInterface(
+            fn=interface,
+            examples=["What is the video about?", "key points of the video"],
+            title=f"Question Mode - {title}",
+        )
+        
+        ui.launch()

requirements.txt

@@ -0,0 +1,11 @@
+torch
+transformers
+nltk
+youtube_transcript_api
+accelerate
+langchain
+yt-dlp
+rich
+chromadb
+InstructorEmbedding
+sentence_transformers

utils.py

@@ -0,0 +1,41 @@
+class LangChainChunker:
+    def __init__(self, text):
+        self.text = text
+    
+    def chunker(self, size=1000):
+        from langchain.text_splitter import CharacterTextSplitter
+        
+        # attach the duration of the video to the chunk
+        # [[chunk, duration]]
+        
+        text_splitter = CharacterTextSplitter(
+            separator=" ",
+            chunk_size=size,
+            chunk_overlap=0.9,
+        )
+        
+        return text_splitter.split_text(self.text)
+    
+    def __sizeof__(self) -> int:
+        count = 0
+        for _ in self.text:
+            count += 1
+        return count
+
+
+def getSubsText(video_id="", getGenerated=False):
+    from youtube_transcript_api import YouTubeTranscriptApi as ytapi
+    from youtube_transcript_api.formatters import TextFormatter
+
+    tList = ytapi.list_transcripts(video_id)
+    data = ""
+    if getGenerated:
+        # TODO: implement getGenerated
+        pass
+    
+    for t in tList:
+        data = t.fetch()
+    
+    return (TextFormatter().format_transcript(data)).replace("\n", " ")
+
+