Update app.py

app.py

@@ -1,161 +1,9 @@
-import os
-import gradio as gr
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
-from langchain_community.llms import HuggingFacePipeline
-from langchain_community.document_loaders import TextLoader, PyPDFLoader
-from langchain.text_splitter import CharacterTextSplitter
-from langchain_community.vectorstores import FAISS
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain.prompts import PromptTemplate
-from langchain.chains import RetrievalQA
-from huggingface_hub import login
-import diarization
-import shutil
-import spaces
-import time
-from langdetect import detect
-
-# Set environment variable to disable tokenizers parallelism warning
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
-
-# 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)
-
-# Language detection function
-def detect_language(text):
-    try:
-        return detect(text)
-    except:
-        return "en"  # default to English if detection fails
-
-# 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 = "mistralai/Mistral-7B-Instruct-v0.3"
-            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_length = 2000,
-                max_new_tokens=512,
-                temperature=0.8,
-            )
-        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=150)
-    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().strip()
-
-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):
-    loader = TextLoader(file_path)
-    documents = loader.load()
-    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
-    texts = text_splitter.split_documents(documents)
-    return texts
-
-attachments_knowledge = load_knowledge("knowledge/bartholomew_attachments_definitions - no int.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_documents(attachments_knowledge, embeddings)
-bigfive_db = FAISS.from_documents(bigfive_knowledge, embeddings)
-personalities_db = FAISS.from_documents(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
-
-    def create_prompt(self, task):
-        return PromptTemplate(
-            template=task + "\n\nContext: {context}\n\nTask: {question}\n\n-----------\n\nAnswer: ",
-            input_variables=["context", "question"]
-        )
-
-    @spaces.GPU(duration=200)
-    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(),
-                chain_type_kwargs={"prompt": self.create_prompt(attachments_task)}
-            )
-            self.bigfive_chain = RetrievalQA.from_chain_type(
-                llm=llm,
-                chain_type="stuff",
-                retriever=bigfive_db.as_retriever(),
-                chain_type_kwargs={"prompt": self.create_prompt(bigfive_task)}
-            )
-            self.personalities_chain = RetrievalQA.from_chain_type(
-                llm=llm,
-                chain_type="stuff",
-                retriever=personalities_db.as_retriever(),
-                chain_type_kwargs={"prompt": self.create_prompt(personalities_task)}
-            )
-        return self.attachments_chain, self.bigfive_chain, self.personalities_chain
-
-lazy_chains = LazyChains(lazy_llm)
-
-@spaces.GPU(duration=150)
-
-def count_words_and_tokens(text):
-    words = len(text.split())
-    tokens = len(AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3").tokenize(text))
-    return words, tokens
-
 
 @spaces.GPU(duration=150)
 def process_input(input_file):
     start_time = time.time()
     
-    yield 0, "Processing file...", None, None, None, None, None, None, None, None
+    yield 0, "Processing file...", None, None, None, None, None, None, None
     
     file_extension = os.path.splitext(input_file.name)[1].lower()
     
@@ -176,7 +24,7 @@ def process_input(input_file):
         temp_video_path = "temp_video" + file_extension
         shutil.copy2(input_file.name, temp_video_path)
         
-        yield 0.2, "Transcribing video...", None, None, None, None, None, None, None, temp_video_path
+        yield 0.2, "Transcribing video...", None, None, None, None, None, None, temp_video_path
         
         language = "en"  # Default to English for video files
         diarization.process_video(temp_video_path, hf_token, language)
@@ -188,24 +36,24 @@ def process_input(input_file):
         input_info = f"Video transcribed. Words: {words}, Tokens: {tokens}"
         video_path = temp_video_path
     else:
-        yield 1, "Unsupported file format. Please upload a TXT, PDF, or video file.", None, None, None, None, None, None, None, None
+        yield 1, "Unsupported file format. Please upload a TXT, PDF, or video file.", None, None, None, None, None, None, None
         return
 
     detected_language = detect_language(content)
     
-    yield 0.4, "Analyzing content...", None, detected_language, input_info, None, None, None, None, video_path
+    yield 0.4, "Analyzing content...", None, detected_language, input_info, None, None, None, video_path
     
     attachments_chain, bigfive_chain, personalities_chain = lazy_chains.get_chains()
     
-    yield 0.6, "Analyzing attachments...", None, detected_language, input_info, None, None, None, None, video_path
+    yield 0.6, "Analyzing attachments...", None, detected_language, input_info, None, None, None, video_path
     attachments_result = attachments_chain({"query": content})
     attachments_answer = attachments_result['result'].split("-----------\n\nAnswer:")[-1].strip()
     
-    yield 0.7, "Analyzing Big Five traits...", None, detected_language, input_info, attachments_answer, None, None, None, video_path
+    yield 0.7, "Analyzing Big Five traits...", None, detected_language, input_info, attachments_answer, None, None, video_path
     bigfive_result = bigfive_chain({"query": content})
     bigfive_answer = bigfive_result['result'].split("-----------\n\nAnswer:")[-1].strip()
     
-    yield 0.8, "Analyzing personalities...", None, detected_language, input_info, attachments_answer, bigfive_answer, None, None, video_path
+    yield 0.8, "Analyzing personalities...", None, detected_language, input_info, attachments_answer, bigfive_answer, None, video_path
     personalities_result = personalities_chain({"query": content})
     personalities_answer = personalities_result['result'].split("-----------\n\nAnswer:")[-1].strip()
     
@@ -225,7 +73,7 @@ def create_interface():
             input_file = gr.File(label="Upload File (TXT, PDF, or Video)")
         
         with gr.Column():
-            progress = gr.Slider(label="Progress", minimum=0, maximum=1, value=0, interactive=False)
+            progress = gr.Progress()
             progress_text = gr.Textbox(label="Status")
             execution_time = gr.Textbox(label="Execution Time", visible=False)
             detected_language = gr.Textbox(label="Detected Language", visible=False)
@@ -237,7 +85,6 @@ def create_interface():
 
         def update_output(progress_value, progress_text, execution_time, detected_lang, input_info, attachments, bigfive, personalities, video):
             return {
-                progress: progress_value,
                 progress_text: progress_text,
                 execution_time: gr.update(value=execution_time, visible=execution_time is not None),
                 detected_language: gr.update(value=detected_lang, visible=detected_lang is not None),
@@ -252,11 +99,10 @@ def create_interface():
             fn=process_input,
             inputs=[input_file],
             outputs=[progress, progress_text, execution_time, detected_language, input_info, attachments_output, bigfive_output, personalities_output, video_output],
-            show_progress=True
         ).then(
             fn=update_output,
             inputs=[progress, progress_text, execution_time, detected_language, input_info, attachments_output, bigfive_output, personalities_output, video_output],
-            outputs=[progress, progress_text, execution_time, detected_language, input_info, attachments_output, bigfive_output, personalities_output, video_output]
+            outputs=[progress_text, execution_time, detected_language, input_info, attachments_output, bigfive_output, personalities_output, video_output]
         )
 
     return iface