Update app.py

app.py

@@ -53,9 +53,9 @@ class LazyPipeline:
                 "text-generation",
                 model=model,
                 tokenizer=tokenizer,
-                max_length = 4000,
+                max_length = 2000,
                 max_new_tokens=512,
-                temperature=0.1,
+                temperature=0.8,
             )
         return self.pipeline
 
@@ -113,7 +113,7 @@ class LazyChains:
 
     def create_prompt(self, task):
         return PromptTemplate(
-            template=task + "\n\nContext: {context}\n\nTask: {question}\n\n---------------------------\n\nAnswer: ",
+            template=task + "\n\nContext: {context}\n\nTask: {question}\n\n-----------\n\nAnswer: ",
             input_variables=["context", "question"]
         )
 
@@ -147,32 +147,38 @@ lazy_chains = LazyChains(lazy_llm)
 
 def count_words_and_tokens(text):
     words = len(text.split())
-    tokens = len(AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3.1-8B-Instruct").tokenize(text))
+    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()
     
+    progress_info = "Processing file..."
+    yield progress_info, None, None, None, None, None, None, None, None
+    
     file_extension = os.path.splitext(input_file.name)[1].lower()
     
     if file_extension == '.txt':
         with open(input_file.name, 'r', encoding='utf-8') as file:
             content = file.read()
         words, tokens = count_words_and_tokens(content)
-        input_info = f"Text file loaded. Words: {words}, Tokens: {tokens}"
+        input_info = f"Text file processed. Words: {words}, Tokens: {tokens}"
         video_path = None
     elif file_extension == '.pdf':
         loader = PyPDFLoader(input_file.name)
         pages = loader.load_and_split()
         content = '\n'.join([page.page_content for page in pages])
         words, tokens = count_words_and_tokens(content)
-        input_info = f"PDF file loaded. Words: {words}, Tokens: {tokens}"
+        input_info = f"PDF file processed. Words: {words}, Tokens: {tokens}"
         video_path = None
     elif file_extension in ['.mp4', '.avi', '.mov']:
         temp_video_path = "temp_video" + file_extension
         shutil.copy2(input_file.name, temp_video_path)
         
+        progress_info = "Transcribing video..."
+        yield progress_info, None, 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)
         
@@ -183,43 +189,61 @@ def process_input(input_file):
         input_info = f"Video transcribed. Words: {words}, Tokens: {tokens}"
         video_path = temp_video_path
     else:
-        return "Unsupported file format. Please upload a TXT, PDF, or video file.", None, None, None, None, None
+        return "Unsupported file format. Please upload a TXT, PDF, or video file.", None, None, None, None, None, None, None, None
 
     detected_language = detect_language(content)
     
+    progress_info = "Analyzing content..."
+    yield progress_info, None, detected_language, input_info, None, None, None, None, video_path
+    
     attachments_chain, bigfive_chain, personalities_chain = lazy_chains.get_chains()
     
     attachments_result = attachments_chain({"query": content})
-    attachments_answer = attachments_result['result'].split("---------------------------\n\nAnswer:")[-1].strip()
+    attachments_answer = attachments_result['result'].split("-----------\n\nAnswer:")[-1].strip()
     
     bigfive_result = bigfive_chain({"query": content})
-    bigfive_answer = bigfive_result['result'].split("---------------------------\n\nAnswer:")[-1].strip()
+    bigfive_answer = bigfive_result['result'].split("-----------\n\nAnswer:")[-1].strip()
     
     personalities_result = personalities_chain({"query": content})
-    personalities_answer = personalities_result['result'].split("---------------------------\n\nAnswer:")[-1].strip()
+    personalities_answer = personalities_result['result'].split("-----------\n\nAnswer:")[-1].strip()
     
     end_time = time.time()
     execution_time = end_time - start_time
     
-    execution_info = f"Execution Time: {execution_time:.2f} seconds\nFile Type: {file_extension}\nDetected Language: {detected_language}"
+    execution_info = f"{execution_time:.2f} seconds"
     
-    return execution_info, input_info, attachments_answer, bigfive_answer, personalities_answer, video_path
-
-iface = gr.Interface(
-    fn=process_input,
-    inputs=gr.File(label="Upload File (TXT, PDF, or Video)"),
-    outputs=[
-        gr.Textbox(label="Execution Information"),
-        gr.Textbox(label="Input Information"),
-        gr.Textbox(label="Attachments Results"),
-        gr.Textbox(label="Big Five Results"),
-        gr.Textbox(label="Personalities Results"),
-        gr.Video(label="Input Video")
-    ],
-    title="Personality Analysis Classification",
-    description="Upload a Video, TXT, or PDF file.",
-    allow_flagging="never"
-)
+    progress_info = "Analysis complete!"
+    
+    yield progress_info, execution_info, detected_language, input_info, attachments_answer, bigfive_answer, personalities_answer, video_path
+
+def create_interface():
+    with gr.Blocks() as iface:
+        gr.Markdown("# Personality Analysis Classification")
+        gr.Markdown("Upload a Video, TXT, or PDF file.")
+        
+        with gr.Row():
+            input_file = gr.File(label="Upload File (TXT, PDF, or Video)")
+        
+        with gr.Column():
+            progress = gr.Textbox(label="Progress")
+            execution_time = gr.Textbox(label="Execution Time", visible=False)
+            detected_language = gr.Textbox(label="Detected Language", visible=False)
+            input_info = gr.Textbox(label="Input Information", visible=False)
+            video_output = gr.Video(label="Input Video", visible=False)
+            attachments_output = gr.Textbox(label="Attachments Results", visible=False)
+            bigfive_output = gr.Textbox(label="Big Five Results", visible=False)
+            personalities_output = gr.Textbox(label="Personalities Results", visible=False)
+
+        input_file.upload(
+            fn=process_input,
+            inputs=[input_file],
+            outputs=[progress, execution_time, detected_language, input_info, attachments_output, bigfive_output, personalities_output, video_output],
+            show_progress=True
+        )
+
+    return iface
+
+iface = create_interface()
 
 # Launch the app
 iface.launch()