Update app.py

app.py

@@ -14,6 +14,12 @@ import shutil
 import spaces
 import time
 from langdetect import detect
+import seaborn as sns
+import matplotlib.pyplot as plt
+import pandas as pd
+import io
+import base64
+import re
 
 # Set environment variable to disable tokenizers parallelism warning
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
@@ -53,11 +59,9 @@ class LazyPipeline:
                 "text-generation",
                 model=model,
                 tokenizer=tokenizer,
-                max_new_tokens=10000,
+                max_new_tokens=4096,
                 do_sample=True,
-                temperature=0.1,
-                top_p=0.95,
-                top_k=50,
+                temperature=0.01,
             )
         return self.pipeline
 
@@ -150,6 +154,33 @@ def count_words_and_tokens(text):
     tokens = len(AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3").tokenize(text))
     return words, tokens
 
+def extract_values(text, keys):
+    values = {}
+    for key in keys:
+        match = re.search(f"{key}:\s*([-]?\d+(?:\.\d+)?)", text)
+        if match:
+            values[key] = float(match.group(1))
+    return values
+
+def create_bar_chart(data, title):
+    plt.figure(figsize=(10, 6))
+    sns.set_style("whitegrid")
+    chart = sns.barplot(x=list(data.keys()), y=list(data.values()), palette="deep")
+    chart.set_xticklabels(chart.get_xticklabels(), rotation=45, ha="right")
+    plt.title(title)
+    plt.tight_layout()
+    
+    # Save the plot to a bytes buffer
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png')
+    buf.seek(0)
+    
+    # Encode the bytes buffer to base64
+    img_str = base64.b64encode(buf.getvalue()).decode()
+    plt.close()
+    
+    return f"data:image/png;base64,{img_str}"
+
 @spaces.GPU(duration=150)
 def process_input(input_file, progress=gr.Progress()):
     start_time = time.time()
@@ -226,34 +257,56 @@ def create_interface():
             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)
-            attachments_output = gr.Textbox(
-                label="Attachments Results",
-                visible=False,
-                lines=10,
-                max_lines=20
-            )
-            bigfive_output = gr.Textbox(
-                label="Big Five Results",
-                visible=False,
-                lines=10,
-                max_lines=20
-            )
-            personalities_output = gr.Textbox(
-                label="Personalities Results",
-                visible=False,
-                lines=10,
-                max_lines=20
-            )
+            
+            attachments_output = gr.Textbox(label="Attachments Results", visible=False, lines=10, max_lines=20)
+            attachments_chart = gr.Image(label="Attachments Chart", visible=False)
+            attachments_explanation = gr.Textbox(label="Attachments Explanation", visible=False)
+            attachments_sentence = gr.Textbox(label="Attachments Sentence", visible=False)
+            
+            bigfive_output = gr.Textbox(label="Big Five Results", visible=False, lines=10, max_lines=20)
+            bigfive_chart = gr.Image(label="Big Five Chart", visible=False)
+            bigfive_explanation = gr.Textbox(label="Big Five Explanation", visible=False)
+            bigfive_sentence = gr.Textbox(label="Big Five Sentence", visible=False)
+            
+            personalities_output = gr.Textbox(label="Personalities Results", visible=False, lines=10, max_lines=20)
+            personalities_chart = gr.Image(label="Personalities Chart", visible=False)
+            personalities_explanation = gr.Textbox(label="Personalities Explanation", visible=False)
+            personalities_sentence = gr.Textbox(label="Personalities Sentence", visible=False)
 
-        def update_visibility(status, exec_time, lang, info, attachments, bigfive, personalities):
+        def update_visibility_and_charts(status, exec_time, lang, info, attachments, bigfive, personalities):
+            attachments_data = extract_values(attachments, ["Secured", "Anxious-Preoccupied", "Dismissive-Avoidant", "Fearful-Avoidant"])
+            attachments_chart_img = create_bar_chart(attachments_data, "Attachment Styles")
+            
+            bigfive_data = extract_values(bigfive, ["Extraversion", "Agreeableness", "Conscientiousness", "Neuroticism", "Openness"])
+            bigfive_chart_img = create_bar_chart(bigfive_data, "Big Five Traits")
+            
+            personalities_data = extract_values(personalities, ["Depressed", "Paranoid", "Schizoid-Schizotypal", "Antisocial-Psychopathic", "Borderline-Dysregulated", "Hysteric-Histrionic", "Narcissistic", "Anxious-Avoidant", "Dependent-Victimized", "Obsessional"])
+            personalities_chart_img = create_bar_chart(personalities_data, "Personality Traits")
+            
+            attachments_explanation = re.search(r"Explanation:(.*?)(?=Sentence:|$)", attachments, re.DOTALL)
+            attachments_sentence = re.search(r"Sentence:(.*?)$", attachments, re.DOTALL)
+            
+            bigfive_explanation = re.search(r"Explanation:(.*?)(?=Sentence:|$)", bigfive, re.DOTALL)
+            bigfive_sentence = re.search(r"Sentence:(.*?)$", bigfive, re.DOTALL)
+            
+            personalities_explanation = re.search(r"Explanation:(.*?)(?=Sentence:|$)", personalities, re.DOTALL)
+            personalities_sentence = re.search(r"Sentence:(.*?)$", personalities, re.DOTALL)
+            
             return {
                 progress_text: status,
                 execution_time: gr.update(value=exec_time, visible=True),
                 detected_language: gr.update(value=lang, visible=True),
                 input_info: gr.update(value=info, visible=True),
                 attachments_output: gr.update(value=attachments, visible=True),
+                attachments_chart: gr.update(value=attachments_chart_img, visible=True),
+                attachments_explanation: gr.update(value=attachments_explanation.group(1).strip() if attachments_explanation else "", visible=True),
+                attachments_sentence: gr.update(value=attachments_sentence.group(1).strip() if attachments_sentence else "", visible=True),
                 bigfive_output: gr.update(value=bigfive, visible=True),
-                personalities_output: gr.update(value=personalities, visible=True)
+                bigfive_chart: gr.update(value=bigfive_chart_img, visible=True),
+                bigfive_explanation: gr.update(value=bigfive_explanation.group(1).strip() if bigfive_explanation else "", visible=True),
+                bigfive_sentence: gr.update(value=bigfive_sentence.group(1).strip() if bigfive_sentence else "", visible=True),
+                personalities_explanation: gr.update(value=personalities_explanation.group(1).strip() if personalities_explanation else "", visible=True),
+                personalities_sentence: gr.update(value=personalities_sentence.group(1).strip() if personalities_sentence else "", visible=True)
             }
 
         input_file.upload(
@@ -261,9 +314,12 @@ def create_interface():
             inputs=[input_file],
             outputs=[progress_text, execution_time, detected_language, input_info, attachments_output, bigfive_output, personalities_output]
         ).then(
-            fn=update_visibility,
+            fn=update_visibility_and_charts,
             inputs=[progress_text, execution_time, detected_language, input_info, attachments_output, bigfive_output, personalities_output],
-            outputs=[progress_text, execution_time, detected_language, input_info, attachments_output, bigfive_output, personalities_output]
+            outputs=[progress_text, execution_time, detected_language, input_info, 
+                     attachments_output, attachments_chart, attachments_explanation, attachments_sentence,
+                     bigfive_output, bigfive_chart, bigfive_explanation, bigfive_sentence,
+                     personalities_output, personalities_chart, personalities_explanation, personalities_sentence]
         )
 
     return iface