app.py

import os
import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
from langchain.llms import HuggingFacePipeline
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.chains import RetrievalQA
from huggingface_hub import login
import diarization
import shutil
import spaces
import time

# Get Hugging Face token from Space secret
hf_token = os.environ.get('hf_secret')
if not hf_token:
    raise ValueError("HF_TOKEN not found in environment variables. Please set it in the Space secrets.")

# Login to Hugging Face
login(token=hf_token)

# Lazy initialization for the pipeline
class LazyPipeline:
    def __init__(self):
        self.pipeline = None

    @spaces.GPU(duration=250)
    def get_pipeline(self):
        if self.pipeline is None:
            import torch
            model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"
            tokenizer = AutoTokenizer.from_pretrained(model_name)
            model = AutoModelForCausalLM.from_pretrained(
                model_name,
                torch_dtype=torch.float16,
                device_map="auto",
            )
            self.pipeline = pipeline(
                "text-generation",
                model=model,
                tokenizer=tokenizer,
                max_new_tokens=512,
                temperature=0.5,
                top_p=0.95,
                repetition_penalty=1.15
            )
        return self.pipeline

lazy_pipe = LazyPipeline()

# Create a LangChain wrapper around the pipeline
class LazyLLM:
    def __init__(self, lazy_pipeline):
        self.lazy_pipeline = lazy_pipeline
        self.llm = None

    @spaces.GPU(duration=250)
    def get_llm(self):
        if self.llm is None:
            pipe = self.lazy_pipeline.get_pipeline()
            self.llm = HuggingFacePipeline(pipeline=pipe)
        return self.llm

lazy_llm = LazyLLM(lazy_pipe)

# Load instruction files
def load_instructions(file_path):
    with open(file_path, 'r') as file:
        return file.read()

general_task = load_instructions("tasks/general_task.txt")
attachments_task = load_instructions("tasks/Attachments_task.txt")
bigfive_task = load_instructions("tasks/BigFive_task.txt")
personalities_task = load_instructions("tasks/Personalities_task.txt")

# Load knowledge files
def load_knowledge(file_path):
    with open(file_path, 'r') as file:
        return file.read()

attachments_knowledge = load_knowledge("knowledge/bartholomew_attachments_definitions.txt")
bigfive_knowledge = load_knowledge("knowledge/bigfive_definitions.txt")
personalities_knowledge = load_knowledge("knowledge/personalities_definitions.txt")

# Create vector stores
embeddings = HuggingFaceEmbeddings()
attachments_db = FAISS.from_texts([attachments_knowledge], embeddings)
bigfive_db = FAISS.from_texts([bigfive_knowledge], embeddings)
personalities_db = FAISS.from_texts([personalities_knowledge], embeddings)

# Lazy initialization for retrieval chains
class LazyChains:
    def __init__(self, lazy_llm):
        self.lazy_llm = lazy_llm
        self.attachments_chain = None
        self.bigfive_chain = None
        self.personalities_chain = None

    @spaces.GPU(duration=120)
    def get_chains(self):
        if self.attachments_chain is None:
            llm = self.lazy_llm.get_llm()
            self.attachments_chain = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=attachments_db.as_retriever())
            self.bigfive_chain = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=bigfive_db.as_retriever())
            self.personalities_chain = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=personalities_db.as_retriever())
        return self.attachments_chain, self.bigfive_chain, self.personalities_chain

lazy_chains = LazyChains(lazy_llm)

# Function to process video file
@spaces.GPU(duration=120)
def process_video(video_file):
    start_time = time.time()

    # Copy the uploaded video file to a temporary location
    temp_video_path = "temp_video.mp4"
    shutil.copy2(video_file.name, temp_video_path)

    # Initialize progress bar
    progress = gr.Progress()

    # Display progress bar for diarization
    progress(0, desc="Starting Diarization...")
    # Process the video using the diarization script
    language = "en"
    diarization.process_video(temp_video_path, hf_token, language)
    progress(50, desc="Diarization Complete.")

    # The SRT file will be created with the same name as the video file but with .srt extension
    srt_path = temp_video_path.replace(".mp4", "_combined.srt")

    # Read the content of the SRT file
    with open(srt_path, 'r', encoding='utf-8') as file:
        srt_content = file.read()

    # Get the chains
    attachments_chain, bigfive_chain, personalities_chain = lazy_chains.get_chains()

    # Process with LangChain and display progress bars
    progress(50, desc="Processing Attachments Analysis...")
    attachments_result = attachments_chain.run(srt_content)
    progress(70, desc="Attachments Analysis Complete.")

    progress(70, desc="Processing Big Five Analysis...")
    bigfive_result = bigfive_chain.run(srt_content)
    progress(90, desc="Big Five Analysis Complete.")

    progress(90, desc="Processing Personalities Analysis...")
    personalities_result = personalities_chain.run(srt_content)
    progress(100, desc="Personalities Analysis Complete.")

    # Combine results
    final_result = f"Attachments Analysis:\n{attachments_result}\n\nBig Five Analysis:\n{bigfive_result}\n\nPersonalities Analysis:\n{personalities_result}"

    end_time = time.time()
    execution_time = end_time - start_time

    # Only return execution time and final result
    final_result_with_time = f"Execution Time: {execution_time:.2f} seconds\n\n{final_result}"

    return final_result_with_time

# Create Gradio interface
iface = gr.Interface(
    fn=process_video,
    inputs=gr.File(label="Upload Video File"),
    outputs=gr.Textbox(label="Results"),
    title="Video Analysis with Meta-Llama-3.1-8B-Instruct",
    description="Upload a video file to analyze using RAG techniques with Meta-Llama-3.1-8B-Instruct."
)

# Launch the app
iface.launch()