Update ui/dream_interface.py

ui/dream_interface.py

@@ -1,88 +1,284 @@
 import os
 import sys
-import argparse
-import gradio as gr
+import json
+import torch
 import logging
+import gradio as gr
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from datetime import datetime
 
 # Set up logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-# Add the current directory to the path
-current_dir = os.path.dirname(os.path.abspath(__file__))
-sys.path.append(current_dir)
-
-def main():
-    parser = argparse.ArgumentParser(description="Freudian Dream Analyzer Interface")
-    parser.add_argument("--share", action="store_true",
-                      help="Create a shareable link")
-    parser.add_argument("--debug", action="store_true",
-                      help="Run in debug mode")
-    parser.add_argument("--hf_space", action="store_true",
-                      help="Running in Hugging Face Space")
-    args = parser.parse_args()
-
-    # Auto-detect if we're in Hugging Face Space
-    in_hf_space = args.hf_space or os.path.exists("/home/user")
-    if in_hf_space and not args.hf_space:
-        logger.info("Detected Hugging Face Space environment")
-        args.hf_space = True
-
-    # Import the dream interface function
-    try:
-        from ui.dream_interface import create_dream_interface
-        logger.info("Successfully imported create_dream_interface")
-    except ImportError as e:
-        logger.error(f"Error importing create_dream_interface: {e}")
-        # Try to add the ui directory to the path
-        ui_dir = os.path.join(current_dir, "ui")
-        if os.path.exists(ui_dir):
-            sys.path.append(ui_dir)
-            try:
-                from dream_interface import create_dream_interface
-                logger.info("Successfully imported create_dream_interface from ui directory")
-            except ImportError as e2:
-                logger.error(f"Error importing create_dream_interface from ui directory: {e2}")
-                raise
-        else:
+class DreamAnalyzer:
+    def __init__(self, model_path):
+        self.model_path = model_path
+        self.pipe, self.model_format = self._load_model()
+        
+    def _load_model(self):
+        logger.info(f"Loading model from {self.model_path}...")
+        
+        # Check if this is a local path or HF Hub model
+        is_local_path = os.path.exists(self.model_path) if not self.model_path.startswith("http") else False
+        model_format = "llama"  # Default format
+        
+        try:
+            # Check if we're in a CPU-only environment
+            cuda_available = torch.cuda.is_available()
+            logger.info(f"CUDA available: {cuda_available}")
+            
+            # Use a pre-defined model from Hugging Face Hub that can run on CPU
+            logger.info("Loading CPU-compatible model from Hugging Face Hub...")
+            # Use a smaller model like Phi-2 that can run on CPU
+            model_id = "microsoft/phi-2"
+            
+            tokenizer = AutoTokenizer.from_pretrained(model_id)
+            model = AutoModelForCausalLM.from_pretrained(
+                model_id,
+                torch_dtype=torch.float32,  # Use float32 for CPU
+                device_map="auto"
+            )
+            
+            logger.info(f"Special tokens: {tokenizer.special_tokens_map}")
+            model_format = "generic"  # Default format for most models
+            
+            # Create text generation pipeline
+            text_generation_pipeline = pipeline(
+                "text-generation",
+                model=model,
+                tokenizer=tokenizer,
+                max_new_tokens=512,  # Reduced for faster inference
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.9,
+                repetition_penalty=1.2,
+            )
+            logger.info("Text generation pipeline created successfully")
+            
+        except Exception as e:
+            logger.error(f"Error loading model: {e}")
+            # Print more detailed debug information
+            logger.error(f"Directory contents of model path:")
+            if is_local_path:
+                try:
+                    logger.error(os.listdir(self.model_path))
+                except Exception as dir_e:
+                    logger.error(f"Could not list directory: {dir_e}")
             raise
-
-    # Use a dummy path - our new interface ignores it and uses a CPU-compatible model
-    model_path = "dummy_path"
-    logger.info(f"Using dummy model path for CPU-compatible model")
+        
+        return text_generation_pipeline, model_format
     
-    # Create the interface
+    def analyze_dream(self, dream_text):
+        if not dream_text or not dream_text.strip():
+            logger.warning("Empty dream text submitted")
+            return "Please enter a dream to analyze."
+        
+        logger.info(f"Analyzing dream: {dream_text[:100]}...")  # Log just the first 100 chars
+        
+        try:
+            # Create a prompt based on the model format
+            prompt = f"You are a Freudian dream analyst. Analyze this dream from a Freudian perspective:\n\n{dream_text}\n\nProvide an analysis that includes manifest content, latent content, and psychological symbolism."
+            
+            logger.info("Submitting prompt to model...")
+            try:
+                # Generate the analysis
+                result = self.pipe(prompt, return_full_text=False, max_new_tokens=512)
+                logger.info("Model returned a response")
+                
+                if not result or len(result) == 0:
+                    logger.error("Empty result from model")
+                    return "<h3>Analysis Error</h3><p>The model returned an empty response. Please try again with a different dream description.</p>"
+                
+                analysis_text = result[0]["generated_text"]
+                logger.info(f"Analysis text length: {len(analysis_text)}")
+                logger.info(f"First 100 chars of analysis: {analysis_text[:100]}")
+                
+                formatted_analysis = self._format_analysis(analysis_text)
+                logger.info("Successfully formatted the analysis")
+                return formatted_analysis
+                
+            except Exception as e:
+                logger.error(f"Error during model inference: {e}")
+                return f"<h3>Analysis Error</h3><p>Error during dream analysis: {str(e)}</p><p>Please try again with a shorter dream description or contact support.</p>"
+                
+        except Exception as e:
+            logger.error(f"Error analyzing dream: {e}")
+            return f"<h3>Error Analyzing Dream</h3><p>There was a problem analyzing your dream: {str(e)}</p><p>Please try again or contact support if the issue persists.</p>"
+    
+    def _format_analysis(self, analysis_text):
+        try:
+            logger.info("Formatting analysis...")
+            
+            # Attempt to structure the analysis with headings
+            sections = [
+                ("Dream Analysis", ""),
+                ("Manifest Content", ""),
+                ("Potential Freudian Symbols", ""),
+                ("Latent Content Analysis", ""),
+                ("Possible Psychological Mechanisms", "")
+            ]
+            
+            # Start with a formatted version that includes headings
+            formatted = "<h3>Freudian Dream Analysis</h3>"
+            formatted += f"<p>{analysis_text}</p>"
+            
+            # If after all formatting we have no content, provide a default message
+            if not formatted.strip():
+                logger.warning("Empty formatted analysis")
+                return "<h3>Analysis</h3><p>The model did not provide a clear analysis. Please try again with a more detailed dream description.</p>"
+                
+            logger.info("Successfully formatted analysis")
+            return formatted
+        except Exception as e:
+            logger.error(f"Error in _format_analysis: {e}")
+            # Return the original text if formatting fails
+            return f"<h3>Dream Analysis</h3><p>{analysis_text}</p>"
+    
+    def record_feedback(self, dream_text, analysis, rating, improvement_suggestion, feedback_dir="feedback_data"):
+        try:
+            os.makedirs(feedback_dir, exist_ok=True)
+            
+            feedback = {
+                "timestamp": datetime.now().isoformat(),
+                "dream_text": dream_text,
+                "model_analysis": analysis,
+                "rating": rating,
+                "improvement_suggestion": improvement_suggestion
+            }
+            
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            feedback_file = os.path.join(feedback_dir, f"feedback_{timestamp}.json")
+            
+            with open(feedback_file, "w") as f:
+                json.dump(feedback, f, indent=2)
+            
+            master_file = os.path.join(feedback_dir, "all_feedback.json")
+            
+            if os.path.exists(master_file):
+                with open(master_file, "r") as f:
+                    try:
+                        all_feedback = json.load(f)
+                    except json.JSONDecodeError:
+                        all_feedback = []
+            else:
+                all_feedback = []
+            
+            all_feedback.append(feedback)
+            
+            with open(master_file, "w") as f:
+                json.dump(all_feedback, f, indent=2)
+            
+            return "Thank you for your feedback! It will help improve future versions of the dream analyzer."
+        except Exception as e:
+            logger.error(f"Error recording feedback: {e}")
+            return f"Error recording feedback: {str(e)}"
+
+# This is the function that needs to be exported
+def create_dream_interface(model_path):
     try:
-        logger.info("Creating interface...")
-        interface = create_dream_interface(model_path)
-        logger.info("Interface created successfully")
+        logger.info(f"Creating dream interface with model path: {model_path}")
+        analyzer = DreamAnalyzer(model_path)
+        
+        def analyze_dream_fn(dream_text):
+            logger.info(f"Analyze button clicked with dream text: {dream_text[:50]}...")
+            try:
+                result = analyzer.analyze_dream(dream_text)
+                logger.info("Analysis complete, returning result")
+                return result
+            except Exception as e:
+                logger.error(f"Error in analyze_dream_fn: {e}")
+                return f"<h3>Error</h3><p>An error occurred while analyzing your dream: {str(e)}</p>"
+        
+        def submit_feedback(dream_text, analysis, rating, suggestion):
+            try:
+                return analyzer.record_feedback(dream_text, analysis, rating, suggestion)
+            except Exception as e:
+                logger.error(f"Error in submit_feedback: {e}")
+                return f"Error submitting feedback: {str(e)}"
         
-        # Launch with appropriate settings
-        if args.debug:
-            logger.info("Launching in debug mode...")
-            interface.launch(
-                server_name="0.0.0.0", 
-                server_port=7860, 
-                share=args.share, 
-                debug=True,
-                show_error=True
+        with gr.Blocks(title="Freudian Dream Analyzer", theme=gr.themes.Soft()) as interface:
+            gr.Markdown("# Freudian Dream Analyzer")
+            gr.Markdown("""
+            This tool analyzes dreams from a Freudian psychoanalytic perspective. 
+            Enter the description of a dream, and the model will provide an interpretation 
+            based on Freudian theory of manifest content, latent content, and symbolic meanings.
+            
+            *Note: Dream analysis is subjective and this tool provides one possible interpretation 
+            based on Freudian psychoanalysis.*
+            """)
+            
+            # Input section
+            with gr.Row():
+                with gr.Column(scale=2):
+                    dream_input = gr.Textbox(
+                        lines=10, 
+                        placeholder="Enter your dream here...",
+                        label="Dream Description"
+                    )
+                    analyze_btn = gr.Button("Analyze Dream", variant="primary")
+                    
+            # Output section
+            with gr.Row():
+                with gr.Column(scale=2):
+                    analysis_output = gr.HTML(
+                        label="Freudian Analysis",
+                        elem_id="analysis_output"
+                    )
+            
+            # Feedback section
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown("### Help improve the Dream Analyzer")
+                    feedback_rating = gr.Slider(
+                        minimum=1, 
+                        maximum=5, 
+                        step=1, 
+                        value=3, 
+                        label="Rate the quality of this analysis (1-5)"
+                    )
+                    feedback_text = gr.Textbox(
+                        lines=3, 
+                        placeholder="Suggestions for improvement (optional)",
+                        label="Feedback"
+                    )
+                    feedback_btn = gr.Button("Submit Feedback")
+                    feedback_output = gr.Textbox(label="")
+            
+            # Add a simple example for testing
+            example_dreams = [
+                ["I was flying over a vast ocean, feeling both exhilarated and afraid. Suddenly I started falling, but before hitting the water I woke up."],
+                ["I was in my childhood home, but the rooms were bigger and had doors I had never seen before. I opened one door and found a room full of old toys."],
+                ["I was being chased through a dark forest by a figure I couldn't see clearly. I kept running but my legs felt heavy and I couldn't move fast enough."],
+                ["I was taking an important exam but realized I hadn't studied at all. The questions made no sense and the clock was ticking loudly. Everyone else seemed to know the answers."],
+                ["I was at a party with friends when I suddenly realized I wasn't wearing any clothes. Everyone was talking normally as if nothing was wrong, but I felt deeply embarrassed."]
+            ]
+            
+            gr.Examples(
+                examples=example_dreams,
+                inputs=dream_input
             )
-        else:
-            logger.info("Launching in normal mode...")
-            interface.launch(
-                server_name="0.0.0.0", 
-                server_port=7860, 
-                share=args.share,
-                show_error=True
+            
+            # Set up event handlers with progress indicator
+            analyze_btn.click(
+                fn=lambda x: "Analyzing dream... please wait.",
+                inputs=dream_input,
+                outputs=analysis_output,
+                queue=False  # Run immediately without queuing
+            ).then(
+                fn=analyze_dream_fn,
+                inputs=dream_input,
+                outputs=analysis_output,
+                queue=True  # Use queue for the actual analysis
             )
+            
+            feedback_btn.click(
+                fn=submit_feedback,
+                inputs=[dream_input, analysis_output, feedback_rating, feedback_text],
+                outputs=feedback_output
+            )
+            
+        return interface
     except Exception as e:
-        logger.error(f"Error creating or launching interface: {e}")
-        raise
-
-if __name__ == "__main__":
-    try:
-        logger.info("Starting Freudian Dream Analyzer application")
-        main()
-    except Exception as e:
-        logger.error(f"Application crashed: {e}")
+        logger.error(f"Error creating dream interface: {e}")
         raise