Update processing.py

processing.py

@@ -1,233 +1,102 @@
-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
-from langchain_community.llms import HuggingFacePipeline
+from langchain.output_parsers import StructuredOutputParser, ResponseSchema
 from langchain.prompts import PromptTemplate
-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 transcription_diarization import process_video
-from output_parser import get_prompt_template, attachment_parser, bigfive_parser, personality_parser, parse_analysis_output
+from pydantic import BaseModel
+from typing import Dict
 
-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.")
+class AttachmentStyle(BaseModel):
+    speaker: str
+    secured: float
+    anxious_preoccupied: float
+    dismissive_avoidant: float
+    fearful_avoidant: float
+    self_rating: int
+    others_rating: int
+    anxiety: int
+    avoidance: int
+    explanation: str
 
-login(token=hf_token)
+class BigFiveTraits(BaseModel):
+    speaker: str
+    extraversion: int
+    agreeableness: int
+    conscientiousness: int
+    neuroticism: int
+    openness: int
+    explanation: str
 
-def load_instructions(file_path):
-    with open(file_path, 'r') as file:
-        return file.read().strip()
+class PersonalityDisorder(BaseModel):
+    speaker: str
+    depressed: int
+    paranoid: int
+    schizoid_schizotypal: int
+    antisocial_psychopathic: int
+    borderline_dysregulated: int
+    narcissistic: int
+    anxious_avoidant: int
+    dependent_victimized: int
+    obsessional: int
+    explanation: str
 
-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
+attachment_response_schemas = [
+    ResponseSchema(name="speaker", description="The name or number of the speaker"),
+    ResponseSchema(name="secured", description="Probability of secured attachment style (0-1)"),
+    ResponseSchema(name="anxious_preoccupied", description="Probability of anxious-preoccupied attachment style (0-1)"),
+    ResponseSchema(name="dismissive_avoidant", description="Probability of dismissive-avoidant attachment style (0-1)"),
+    ResponseSchema(name="fearful_avoidant", description="Probability of fearful-avoidant attachment style (0-1)"),
+    ResponseSchema(name="self_rating", description="Self rating (0-10)"),
+    ResponseSchema(name="others_rating", description="Others rating (0-10)"),
+    ResponseSchema(name="anxiety", description="Anxiety rating (0-10)"),
+    ResponseSchema(name="avoidance", description="Avoidance rating (0-10)"),
+    ResponseSchema(name="explanation", description="Brief explanation of the attachment style")
+]
 
-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")
+bigfive_response_schemas = [
+    ResponseSchema(name="speaker", description="The name or number of the speaker"),
+    ResponseSchema(name="extraversion", description="Extraversion rating (-10 to 10)"),
+    ResponseSchema(name="agreeableness", description="Agreeableness rating (-10 to 10)"),
+    ResponseSchema(name="conscientiousness", description="Conscientiousness rating (-10 to 10)"),
+    ResponseSchema(name="neuroticism", description="Neuroticism rating (-10 to 10)"),
+    ResponseSchema(name="openness", description="Openness rating (-10 to 10)"),
+    ResponseSchema(name="explanation", description="Brief explanation of the Big Five traits")
+]
 
-embeddings = HuggingFaceEmbeddings()
-attachments_db = FAISS.from_documents(load_knowledge("knowledge/bartholomew_attachments_definitions_no_items_no_in.txt"), embeddings)
-bigfive_db = FAISS.from_documents(load_knowledge("knowledge/bigfive_definitions_no_items.txt"), embeddings)
-personalities_db = FAISS.from_documents(load_knowledge("knowledge/personalities_definitions.txt"), embeddings)
+personality_response_schemas = [
+    ResponseSchema(name="speaker", description="The name or number of the speaker"),
+    ResponseSchema(name="depressed", description="Depressed rating (0-4)"),
+    ResponseSchema(name="paranoid", description="Paranoid rating (0-4)"),
+    ResponseSchema(name="schizoid_schizotypal", description="Schizoid-Schizotypal rating (0-4)"),
+    ResponseSchema(name="antisocial_psychopathic", description="Antisocial-Psychopathic rating (0-4)"),
+    ResponseSchema(name="borderline_dysregulated", description="Borderline-Dysregulated rating (0-4)"),
+    ResponseSchema(name="narcissistic", description="Narcissistic rating (0-4)"),
+    ResponseSchema(name="anxious_avoidant", description="Anxious-Avoidant rating (0-4)"),
+    ResponseSchema(name="dependent_victimized", description="Dependent-Victimized rating (0-4)"),
+    ResponseSchema(name="obsessional", description="Obsessional rating (0-4)"),
+    ResponseSchema(name="explanation", description="Brief explanation of the personality disorders")
+]
 
-def detect_language(text):
-    try:
-        return detect(text)
-    except:
-        return "en"
+attachment_parser = StructuredOutputParser.from_response_schemas(attachment_response_schemas)
+bigfive_parser = StructuredOutputParser.from_response_schemas(bigfive_response_schemas)
+personality_parser = StructuredOutputParser.from_response_schemas(personality_response_schemas)
 
-class SequentialAnalyzer:
-    def __init__(self, hf_token, seed=42):
-        self.hf_token = hf_token
-        self.model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"
-        self.set_seed(seed)
-        self.model = self.load_model()
-        self.pipe = self.create_pipeline(self.model)
-
-    def set_seed(self, seed):
-        random.seed(seed)
-        np.random.seed(seed)
-        torch.manual_seed(seed)
-        if torch.cuda.is_available():
-            torch.cuda.manual_seed_all(seed)
-
-    def load_model(self):
-        model = AutoModelForCausalLM.from_pretrained(
-            self.model_name,
-            torch_dtype=torch.bfloat16,
-            device_map="auto",
-            use_auth_token=self.hf_token,
-            use_cache=False,
-            load_in_4bit=False
-        )
-        return model
-
-    def create_pipeline(self, model):
-        from transformers import pipeline
-        tokenizer = AutoTokenizer.from_pretrained(self.model_name, use_auth_token=self.hf_token)
-        return pipeline(
-            "text-generation",
-            model=model,
-            top_k=50,
-            top_p=0.8,
-            tokenizer=tokenizer,
-            max_new_tokens=512,
-            temperature=0.3,
-            repetition_penalty=1.2,
-            do_sample=False,
-            truncation=True,
-            bad_words_ids=[[tokenizer.encode(char, add_special_tokens=False)[0]] for char in "*"]
-        )
-        
-    def post_process_output(self, output):
-        return re.sub(r'[*]', '', output).strip()
-
-    def analyze_task(self, content, task, knowledge_db, analysis_type):
-        tokenizer = AutoTokenizer.from_pretrained(self.model_name, use_auth_token=self.hf_token)
-        
-        input_tokens = len(tokenizer.encode(content))
-        
-        max_input_length = 800
-        encoded_input = tokenizer.encode(content, truncation=True, max_length=max_input_length)
-        truncated_content = tokenizer.decode(encoded_input)
-        
-        if len(encoded_input) == max_input_length:
-            print(f"Warning: Input was truncated from {input_tokens} to {max_input_length} tokens.")
-        
-        llm = HuggingFacePipeline(pipeline=self.pipe)
-    
-        if analysis_type == "attachments":
-            parser = attachment_parser
-        elif analysis_type == "bigfive":
-            parser = bigfive_parser
-        elif analysis_type == "personalities":
-            parser = personality_parser
-        else:
-            raise ValueError(f"Unknown analysis type: {analysis_type}")
-    
-        prompt = get_prompt_template(task, parser)
-        
-        if knowledge_db is None:
-            chain = prompt | llm
-            result = chain.invoke({"text": truncated_content})
-            output = result
-        else:
-            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" + parser.get_format_instructions() + "\n\nAnalysis:",
-                        input_variables=["context", "question"]
-                    )
-                }
-            )
-            result = chain.run(truncated_content)
-            output = result
-    
-        print(f"Raw model output: {output}")
-    
-        try:
-            cleaned_output = self.post_process_output(output)
-            parsed_output = parser.parse(cleaned_output)
-        except Exception as e:
-            raise ValueError(f"Error parsing output: {e}")
-    
-        # Check if all required keys are present
-        required_keys = {schema.name for schema in parser.response_schemas}
-        missing_keys = required_keys - parsed_output.keys()
-    
-        if missing_keys:
-            raise ValueError(f"Missing some input keys: {missing_keys}")
-    
-        return cleaned_output, input_tokens
+def get_prompt_template(task: str, parser: StructuredOutputParser) -> PromptTemplate:
+    return PromptTemplate(
+        template="Analyze the following text according to the given task:\n\n{task}\n\n{format_instructions}\n\nText: {text}\n\nAnalysis:",
+        input_variables=["text"],
+        partial_variables={
+            "task": task,
+            "format_instructions": parser.get_format_instructions()
+        }
+    )
 
-def process_input(input_file, max_speakers, progress=None):
-    start_time = time.time()
-    
-    def safe_progress(value, desc=""):
-        if progress is not None:
-            try:
-                progress(value, desc=desc)
-            except Exception as e:
-                print(f"Progress update failed: {e}")
-    
-    safe_progress(0, desc="Processing file")
-    
-    if isinstance(input_file, str):
-        file_path = input_file
+def parse_analysis_output(output: str, analysis_type: str) -> Dict[str, BaseModel]:
+    if analysis_type == "attachments":
+        parsed = attachment_parser.parse(output)
+        return {parsed['speaker']: AttachmentStyle(**parsed)}
+    elif analysis_type == "bigfive":
+        parsed = bigfive_parser.parse(output)
+        return {parsed['speaker']: BigFiveTraits(**parsed)}
+    elif analysis_type == "personalities":
+        parsed = personality_parser.parse(output)
+        return {parsed['speaker']: PersonalityDisorder(**parsed)}
     else:
-        file_path = input_file.name
-    
-    file_extension = os.path.splitext(file_path)[1].lower()
-    
-    if file_extension in ['.txt', '.srt']:
-        with open(file_path, 'r', encoding='utf-8') as file:
-            content = file.read()
-        transcription = content
-    elif file_extension == '.pdf':
-        loader = PyPDFLoader(file_path)
-        pages = loader.load_and_split()
-        content = '\n'.join([page.page_content for page in pages])
-        transcription = content
-    elif file_extension in ['.mp4', '.avi', '.mov']:
-        safe_progress(0.2, desc="Processing video...")
-        srt_path = process_video(file_path, hf_token, "en", max_speakers)
-        with open(srt_path, 'r', encoding='utf-8') as file:
-            content = file.read()
-        transcription = content
-        os.remove(srt_path)
-    else:
-        return "Unsupported file format. Please upload a TXT, SRT, PDF, or video file.", None, None, None, None, None, None
-    
-    detected_language = detect_language(content)
-    
-    safe_progress(0.2, desc="Initializing analyzer")
-    analyzer = SequentialAnalyzer(hf_token)
-    
-    tasks = [
-        ("General + Attachments", general_task + "\n\n" + attachments_task, attachments_db, "attachments"),
-        ("General + Big Five", general_task + "\n\n" + bigfive_task, bigfive_db, "bigfive"),
-        ("General + Personalities", general_task + "\n\n" + personalities_task, personalities_db, "personalities")
-    ]
-    
-    results = []
-    tokens = []
-    
-    for i, (task_name, task, db, analysis_type) in enumerate(tasks):
-        safe_progress((i + 1) * 0.2, desc=f"Analyzing {task_name}")
-        answer, task_tokens = analyzer.analyze_task(content, task, db, analysis_type)
-        results.append((answer, analysis_type))
-        tokens.append(task_tokens)
-    
-    end_time = time.time()
-    execution_time = end_time - start_time
-    
-    safe_progress(1.0, desc="Analysis complete!")
-
-    parsed_results = [parse_analysis_output(result, analysis_type) for result, analysis_type in results]
-
-    return (
-        "Analysis complete!",
-        f"{execution_time:.2f} seconds",
-        detected_language,
-        parsed_results[0],  # attachments
-        parsed_results[1],  # bigfive
-        parsed_results[2],  # personalities,
-        transcription
-    )
+        raise ValueError(f"Unknown analysis type: {analysis_type}")