app.py

from stable_whisper import modify_model,results_to_word_srt, results_to_sentence_srt
import whisper
import pysrt
import re
import os
from copy import deepcopy
from typing import List
import os
from langchain import HuggingFaceHub, PromptTemplate, LLMChain
from langchain.document_loaders import UnstructuredPDFLoader
from nltk.corpus import stopwords
from nltk.tokenize import word_tokenize
from nltk import FreqDist
from nltk.metrics import jaccard_distance

from moviepy.video.io.VideoFileClip import VideoFileClip
from moviepy.video.VideoClip import ImageClip
from datetime import datetime
import gradio as gr
import nltk
nltk.download('stopwords')
nltk.download('punkt')

huggingfacehub_api_token = os.getenv("HF_TOKEN")

class VideoQA:
    def __init__(self):
        # self.loader = UnstructuredPDFLoader("/content/Document_ Introduction to Python (1).pdf")
        # self.extracted_text=self.loader.load()

        # self.huggingfacehub_api_token =  # Replace with your Hugging Face token
        self.repo_id = "mistralai/Mistral-7B-Instruct-v0.1"
        self.llm = HuggingFaceHub(
            huggingfacehub_api_token=huggingfacehub_api_token,
            repo_id=self.repo_id,
            model_kwargs={"temperature": 0.2, "max_new_tokens": 800}
        )

    def load_model(self,model_selected):
      """
      Load a pre-trained machine learning model specified by the `model_selected` parameter
      using the `whisper` library and modify it to output word timestamps.

      Parameters:
      -----------
      model_selected : str
          A string specifying the name of the pre-trained machine learning model to load.

      Returns:
      --------
      model : object
          A modified version of the loaded pre-trained machine learning model that outputs
          timestamps for individual words.

      """
      model = whisper.load_model(model_selected)
      modify_model(model)
      return model

    def whisper_result_to_srt(self,result):
        """
        Convert the output of the Whisper speech recognition model into SubRip subtitle format.

        Parameters:
        -----------
        result : dict
            A dictionary containing the output of the Whisper speech recognition model, including word-level
            timestamps.

        Returns:
        --------
        srt : str
            A string in SubRip subtitle format, containing the word-level transcriptions and timing information
            from the Whisper output.

        Notes:
        ------
        This function takes the output of the Whisper speech recognition model, which includes word-level timestamps
        for each segment of the input audio file, and converts it into SubRip subtitle format. The resulting subtitle
        file can be used to display captions or transcripts alongside a video recording of the original audio.
        """
        text = []
        for i,s in enumerate(result['segments']):
            text.append(str(i+1))
            time_start = s['start']
            hours, minutes, seconds = int(time_start/3600), (time_start/60) % 60, (time_start) % 60
            timestamp_start = "%02d:%02d:%06.3f" % (hours, minutes, seconds)
            timestamp_start = timestamp_start.replace('.',',')
            time_end = s['end']
            hours, minutes, seconds = int(time_end/3600), (time_end/60) % 60, (time_end) % 60
            timestamp_end = "%02d:%02d:%06.3f" % (hours, minutes, seconds)
            timestamp_end = timestamp_end.replace('.',',')
            text.append(timestamp_start + " --> " + timestamp_end)
            text.append(s['text'].strip() + "\n")
        return "\n".join(text)

    # model_selected  = 'tiny'
    def transcribe_video(self,vid, model_selected):
        """
        Transcribe the audio in a video file using a pre-trained machine learning model and return the transcription
        and its corresponding timestamps in a subtitle format.

        Parameters:
        -----------
        vid : str
            A string specifying the path to the video file to be transcribed.
        model_selected : str
            A string specifying the name of the pre-trained machine learning model to use for transcription.

        Returns:
        --------
        result : dict
            A dictionary containing the transcription and its corresponding timestamps in a subtitle format.

        """
        model = self.load_model(model_selected)
        options = whisper.DecodingOptions(fp16=False)
        result = model.transcribe(vid, **options.__dict__)
        result['srt'] = self.whisper_result_to_srt(result)
        return result


    def to_srt(self,lines: List[dict], strip=False) -> str:
        """
        lines: List[dict]
            [{start:<start-timestamp-of-text>, end:<end-timestamp-of-text>, text:<str-of-text>}, ...]
        """

        def secs_to_hhmmss(secs):
            mm, ss = divmod(secs, 60)
            hh, mm = divmod(mm, 60)
            return f'{hh:0>2.0f}:{mm:0>2.0f}:{ss:0>6.3f}'.replace(".", ",")

        srt_str = '\n'.join(
            f'{i}\n'
            f'{secs_to_hhmmss(sub["start"])} --> {secs_to_hhmmss(sub["end"])}\n'
            f'{sub["text"].strip() if strip else sub["text"]}\n'
            for i, sub in enumerate(lines, 1))

        # if save_path:
        #     with open(save_path, 'w', encoding='utf-8') as f:
        #         f.write(srt_str)
        #     print(f'Saved: {os.path.abspath(save_path)}')

        return srt_str

    def tighten_timestamps(self,res: dict, end_at_last_word=True, end_before_period=False, start_at_first_word=False) -> dict:
        res = deepcopy(res)
        for i in range(len(res['segments'])):
            if start_at_first_word:
                res['segments'][i]['start'] = res['segments'][i]['word_timestamps'][0]['timestamp']
            if end_before_period and \
                    res['segments'][i]['word_timestamps'][-1] == '.' and \
                    len(res['segments'][i]['word_timestamps']) > 1:
                res['segments'][i]['end'] = res['segments'][i]['word_timestamps'][-2]['timestamp']
            elif end_at_last_word:
                res['segments'][i]['end'] = res['segments'][i]['word_timestamps'][-1]['timestamp']

        return res

    def results_to_sentence_srt(self,res: dict,
                                end_at_last_word=False,
                                end_before_period=False,
                                start_at_first_word=False,
                                strip=False):
        """

        Parameters
        ----------
        res: dict
            results from modified model
        srt_path: str
            output path of srt
        end_at_last_word: bool
            set end-of-sentence to timestamp-of-last-token
        end_before_period: bool
            set end-of-sentence to timestamp-of-last-non-period-token
        start_at_first_word: bool
            set start-of-sentence to timestamp-of-first-token
        strip: bool
            perform strip() on each sentence

        """
        strict = any((end_at_last_word, end_before_period, start_at_first_word))
        segs = self.tighten_timestamps(res,
                                  end_at_last_word=end_at_last_word,
                                  end_before_period=end_before_period,
                                  start_at_first_word=start_at_first_word)['segments'] \
            if strict else res['segments']

        max_idx = len(segs) - 1
        i = 1
        while i <= max_idx:
            if not (segs[i]['end'] - segs[i]['start']):
                if segs[i - 1]['end'] == segs[i]['end']:
                    segs[i - 1]['text'] += (' ' + segs[i]['text'].strip())
                    del segs[i]
                    max_idx -= 1
                    continue
                else:
                    segs[i]['start'] = segs[i - 1]['end']
            i += 1

        srt = self.to_srt(segs, strip=strip)
        return srt



    def extract_timestamps_and_text(self,input_text):
        timestamp_pattern = re.compile(r'(\d{2}:\d{2}:\d{2}.\d{3})\s*-->\s*(\d{2}:\d{2}:\d{2}.\d{3})\n(.+)')

        matches = timestamp_pattern.findall(input_text)

        data = []

        for match in matches:
            start_timestamp, end_timestamp, text = match
            data.append({
                'start_timestamp': start_timestamp,
                'end_timestamp': end_timestamp,
                'text': text.strip()
            })

        return data


    def generate_contract(self,text,question):


        template = """you are the german language and universal language expert .your task is  analyze the given  text and user ask any question about given text answer to the user question.your returning answer must in user's language.otherwise reply i don't know.
        extracted_text:{text}
        user_question:{question}"""

        prompt = PromptTemplate(template=template, input_variables=["text","question"])
        llm_chain = LLMChain(prompt=prompt, verbose=True, llm=self.llm)

        result = llm_chain.run({"text":text,"question":question})
        print()
        print()
        print("this is answer:",result)
        return result


    def preprocess_sentence(self,sentence):
        stop_words = set(stopwords.words('english'))
        words = word_tokenize(sentence.lower())
        filtered_words = [word for word in words if word.isalnum() and word not in stop_words]
        return filtered_words

    def compute_similarity(self,sentence1, sentence2):
        words1 = self.preprocess_sentence(sentence1)
        words2 = self.preprocess_sentence(sentence2)

        freq_dist1 = FreqDist(words1)
        freq_dist2 = FreqDist(words2)

        jaccard = 1 - jaccard_distance(set(freq_dist1), set(freq_dist2))

        return jaccard

    def find_most_similar(self,sentence_list, target_sentence):
        similarities = [self.compute_similarity(target_sentence, sentence) for sentence in sentence_list]

        # Find the index of the most similar sentence
        most_similar_index = similarities.index(max(similarities))

        # Return the most similar sentence
        return sentence_list[most_similar_index]


    def start_end_timestamp(self,result,answer):
      appended_text = []

      for item in result:
          appended_text.append(item['text'])

      # Find the most similar sentence
      matched_sentence = self.find_most_similar(appended_text, answer)
      start_time=""
      end_time=""
      for entry in result:
        if matched_sentence in entry['text']:
          start_time = entry['start_timestamp']
          end_time = entry['end_timestamp']
          print(start_time+"\n"+end_time)
      return start_time,end_time



    def timestamp_to_seconds(self,timestamp):
        time_format = "%H:%M:%S,%f"
        dt = datetime.strptime(timestamp, time_format)
        return dt.hour * 3600 + dt.minute * 60 + dt.second + dt.microsecond / 1e6

    def cut_video(self,input_file, output_file, start_timestamp, end_timestamp):
        # Convert timestamps to seconds
        start_time = self.timestamp_to_seconds(start_timestamp)
        end_time = self.timestamp_to_seconds(end_timestamp)

        # Use moviepy to cut both video and audio
        video_clip = VideoFileClip(input_file).subclip(start_time, end_time)
        video_clip.write_videofile(output_file, codec='libx264', audio_codec='aac', temp_audiofile='temp-audio.m4a', remove_temp=True)

    def main(self,input_video_path,question):

      subtitle = self.transcribe_video(input_video_path,'medium')
      text = subtitle['text']
      answer = self.generate_contract(text,question)

      subrip_text  = self.results_to_sentence_srt(subtitle)
      result = self.extract_timestamps_and_text(subrip_text)
      start_time,end_time = self.start_end_timestamp(result,answer)
      output_video_path = 'output_video.mp4'

      self.cut_video(input_video_path, output_video_path, start_time, end_time)
      return output_video_path

    def gradio_interface(self):

        with gr.Blocks() as demo:
          gr.HTML("""<center><h1>Video Question Answering</h1></center>""")
          with gr.Row():
            video = gr.Video()
          with gr.Row():
            query = gr.Textbox("Query")
          with gr.Row():
            output_video = gr.Video()

          query.submit(self.main,[video,query],output_video)
        demo.launch(debug=True)

if __name__=="__main__":
  video_qa = VideoQA()
  video_qa.gradio_interface()