Update processing.py

processing.py

@@ -1,8 +1,10 @@
 import os
 import time
 import re
+import numpy as np
 from huggingface_hub import login
 import torch
+import random
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from langdetect import detect
 from langchain.chains import RetrievalQA
@@ -28,6 +30,7 @@ def load_instructions(file_path):
     with open(file_path, 'r') as file:
         return file.read().strip()
 
+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")
@@ -74,7 +77,6 @@ class SequentialAnalyzer:
             use_cache=False,
             load_in_4bit=False
         )
-        model.gradient_checkpointing_enable()
         return model
 
     def create_pipeline(self, model):
@@ -109,18 +111,30 @@ class SequentialAnalyzer:
             print(f"Warning: Input was truncated from {input_tokens} to {max_input_length} tokens.")
         
         llm = HuggingFacePipeline(pipeline=self.pipe)
-        chain = RetrievalQA.from_chain_type(
-            llm=llm,
-            chain_type="stuff",
-            retriever=knowledge_db.as_retriever(),
-            chain_type_kwargs={"prompt": PromptTemplate(
-                template=task + "\n\n{context}\n\n{question}\n\n-----------\n\nAnswer: ",
-                input_variables=["context", "question"]
-            )}
-        )
-        
-        result = chain({"query": truncated_content})
-        output = result['result'].split("-----------\n\nAnswer:")[-1].strip()
+
+        if knowledge_db is None:
+            # For general task without specific knowledge DB
+            prompt = PromptTemplate(
+                template=task + "\n\n{question}\n\n-----------\n\nAnswer: ",
+                input_variables=["question"]
+            )
+            chain = prompt | llm
+            result = chain.invoke({"question": truncated_content})
+            output = result.split("-----------\n\nAnswer:")[-1].strip()
+        else:
+            # For tasks with specific knowledge DB
+            chain = RetrievalQA.from_chain_type(
+                llm=llm,
+                chain_type="stuff",
+                retriever=knowledge_db.as_retriever(),
+                chain_type_kwargs={"prompt": PromptTemplate(
+                    template=task + "\n\n{context}\n\n{question}\n\n-----------\n\nAnswer: ",
+                    input_variables=["context", "question"]
+                )}
+            )
+            result = chain({"query": truncated_content})
+            output = result['result'].split("-----------\n\nAnswer:")[-1].strip()
+
         cleaned_output = self.post_process_output(output)
         
         return cleaned_output, input_tokens
@@ -161,7 +175,7 @@ def process_input(input_file, progress=None):
         transcription = content  # Store the transcription
         os.remove(srt_path)
     else:
-        return "Unsupported file format. Please upload a TXT, SRT, PDF, or video file.", None, None, None, None, None, None, None, None, None, None
+        return "Unsupported file format. Please upload a TXT, SRT, PDF, or video file.", None, None, None, None, None, None, None, None, None, None, None, None
         
     detected_language = detect_language(content)
     
@@ -172,6 +186,11 @@ def process_input(input_file, progress=None):
     
     analyzer = SequentialAnalyzer(hf_token)
     
+    safe_progress(0.3, desc="Performing general analysis...")
+    general_answer, general_tokens = analyzer.analyze_task(content, general_task, None)
+    print("General output:\n", general_answer)
+    print(f"General input tokens (before truncation): {general_tokens}")
+
     safe_progress(0.5, desc="Analyzing attachments...")
     attachments_answer, attachments_tokens = analyzer.analyze_task(content, attachments_task, attachments_db)
     print("Attachments output:\n", attachments_answer)
@@ -195,5 +214,5 @@ def process_input(input_file, progress=None):
     safe_progress(1.0, desc="Analysis complete!")
 
     return ("Analysis complete!", execution_info, detected_language, 
-            attachments_answer, bigfive_answer, personalities_answer,
-            original_tokens, attachments_tokens, bigfive_tokens, personalities_tokens, transcription)
+            general_answer, attachments_answer, bigfive_answer, personalities_answer,
+            original_tokens, general_tokens, attachments_tokens, bigfive_tokens, personalities_tokens, transcription)