Update app.py

app.py

@@ -1,6 +1,12 @@
+# Required imports - ensure all dependencies are installed
 import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+try:
+    import bitsandbytes as bnb
+except ImportError:
+    print("WARNING: bitsandbytes not installed. Required for 4-bit quantization.")
+
 import time
 import os
 import json
@@ -52,38 +58,55 @@ def load_model_in_background():
     global MODEL, TOKENIZER, PIPE, MODEL_LOADING, MODEL_LOADED
     try:
         MODEL_LOADING = True
-        print("Starting model loading...")
+        print("Starting model loading process...")
         
-        # Model identifier
+        # Model identifier - using quantized 4-bit version for reduced memory
         model_id = "mistralai/Mistral-7B-Instruct-v0.3"
         
         print("Loading tokenizer...")
-        TOKENIZER = AutoTokenizer.from_pretrained(model_id)
+        # Set tokenizer to use legacy format to avoid issues
+        TOKENIZER = AutoTokenizer.from_pretrained(model_id, legacy_format=True)
         
-        print("Loading model (this may take several minutes)...")
+        print("Loading model with optimized settings for limited memory...")
+        # Configure model loading with 4-bit quantization for minimum memory usage
         MODEL = AutoModelForCausalLM.from_pretrained(
             model_id,
-            torch_dtype=torch.float16,
-            device_map="auto",
-            low_cpu_mem_usage=True
+            torch_dtype=torch.float16,  # Use half precision
+            device_map="auto",  # Automatically distribute across available GPUs/CPU
+            low_cpu_mem_usage=True,
+            load_in_4bit=True,  # Enable 4-bit quantization
+            max_memory={0: "8GiB"},  # Limit memory usage per GPU
+            offload_folder="offload_folder",  # Use disk offloading if needed
+            offload_state_dict=True  # Offload state dict to CPU when possible
         )
         
-        print("Creating pipeline...")
+        print("Creating optimized pipeline...")
+        # Create text generation pipeline with more efficient settings
         PIPE = pipeline(
             "text-generation",
             model=MODEL,
             tokenizer=TOKENIZER,
-            return_full_text=False
+            return_full_text=False,
+            batch_size=1  # Process one batch at a time to reduce memory
         )
         
+        print("Model loading complete!")
         MODEL_LOADING = False
         MODEL_LOADED = True
-        print("Model loaded successfully!")
-        return "Model loaded successfully!"
+        return "Model loaded successfully! Ready to generate responses."
+    except torch.cuda.OutOfMemoryError as e:
+        MODEL_LOADING = False
+        print(f"CUDA out of memory error: {str(e)}")
+        return f"GPU memory error: {str(e)}. Try restarting or using a machine with more GPU memory."
+    except ImportError as e:
+        MODEL_LOADING = False
+        print(f"Import error - missing dependencies: {str(e)}")
+        return f"Missing dependencies: {str(e)}. Try 'pip install -U bitsandbytes transformers accelerate'"
     except Exception as e:
         MODEL_LOADING = False
         print(f"Error loading model: {str(e)}")
-        return f"Error loading model: {str(e)}"
+        error_type = type(e).__name__
+        return f"Error loading model ({error_type}): {str(e)}"
 
 # Function to generate response using the model
 def generate_response(prompt, chat_history, progress=gr.Progress()):
@@ -533,6 +556,9 @@ def create_gradio_interface():
     with gr.Blocks(css=css) as app:
         gr.Markdown("# 🤖 Advanced Mistral-7B-Instruct Chatbot for n8n JSON Generation")
         
+        # Add fallback mode when model loading fails
+        fallback_mode = gr.State(False)
+        
         with gr.Tab("Chat"):
             with gr.Row():
                 with gr.Column(scale=3):
@@ -554,6 +580,49 @@ def create_gradio_interface():
                         )
                         send_btn = gr.Button("Send", scale=1, variant="primary")
                     
+                    # Fallback response mode when model fails to load
+                    def simple_fallback_response(message, history):
+                        # Sample n8n JSON structure for common requests
+                        if any(keyword in message.lower() for keyword in ['json', 'n8n', 'workflow']):
+                            response = """Here's a basic n8n workflow JSON structure you can use:
+
+```json
+{
+  "name": "Simple Workflow",
+  "nodes": [
+    {
+      "parameters": {
+        "values": {
+          "string": [
+            {
+              "name": "data",
+              "value": "Your data here"
+            }
+          ]
+        }
+      },
+      "id": "1",
+      "name": "Set",
+      "type": "n8n-nodes-base.set",
+      "typeVersion": 1,
+      "position": [250, 300]
+    }
+  ],
+  "connections": {},
+  "active": false,
+  "settings": {},
+  "version": 1
+}
+```
+
+You can customize this template with your specific data. If you need a more complex structure or specific node types, please let me know."""
+                        elif 'file' in message.lower() or 'data' in message.lower():
+                            response = "To analyze files or data, please upload your file in the 'File Analysis & JSON Conversion' tab. I'll be able to help you convert it to n8n compatible JSON format."
+                        else:
+                            response = "I'm currently operating in fallback mode because the Mistral-7B model couldn't be loaded. I can still help with basic n8n JSON structures. Try asking for a specific workflow type or JSON structure you need for n8n."
+                        
+                        return history + [(message, response)]
+                    
                     with gr.Row():
                         chat_selector = gr.Dropdown(
                             choices=get_available_chats(),
@@ -574,9 +643,20 @@ def create_gradio_interface():
                 
                 with gr.Column(scale=1):
                     # Model Loading and Settings
-                    load_model_btn = gr.Button("Load Mistral-7B Model", variant="primary")
+                    with gr.Row():
+                        load_model_btn = gr.Button("Load Mistral-7B Model", variant="primary")
+                        use_fallback_btn = gr.Button("Use Simple JSON Mode", variant="secondary")
+                    
                     model_status = gr.Textbox(label="Model Status", value="Not loaded", interactive=False)
                     
+                    # Function to toggle fallback mode
+                    def toggle_fallback_mode(state):
+                        global MODEL_LOADED
+                        if MODEL_LOADED:
+                            return state, "Model is already loaded. No need for fallback mode."
+                        else:
+                            return not state, "Using simple JSON generation mode. Limited functionality but no model loading required."
+                    
                     # System Prompt
                     system_prompt_input = gr.Textbox(
                         label="System Prompt", 
@@ -675,6 +755,76 @@ def create_gradio_interface():
             }
             ```
             
+            ## Common n8n Node Types
+            
+            ### HTTP Request Node
+            ```json
+            {
+              "parameters": {
+                "url": "https://api.example.com/data",
+                "method": "GET",
+                "authentication": "none",
+                "sendHeaders": true,
+                "headerParameters": {
+                  "parameters": [
+                    {
+                      "name": "Content-Type",
+                      "value": "application/json"
+                    }
+                  ]
+                }
+              },
+              "name": "HTTP Request",
+              "type": "n8n-nodes-base.httpRequest",
+              "typeVersion": 1,
+              "position": [250, 300],
+              "id": "1"
+            }
+            ```
+            
+            ### Function Node (JavaScript)
+            ```json
+            {
+              "parameters": {
+                "functionCode": "// Code here\nreturn items;"
+              },
+              "name": "Function",
+              "type": "n8n-nodes-base.function",
+              "typeVersion": 1,
+              "position": [450, 300],
+              "id": "2"
+            }
+            ```
+            
+            ### Set Node (Manual Data)
+            ```json
+            {
+              "parameters": {
+                "values": {
+                  "string": [
+                    {
+                      "name": "fieldName",
+                      "value": "value"
+                    }
+                  ],
+                  "number": [
+                    {
+                      "name": "count",
+                      "value": 42
+                    }
+                  ]
+                }
+              },
+              "name": "Set",
+              "type": "n8n-nodes-base.set",
+              "typeVersion": 1,
+              "position": [650, 300],
+              "id": "3"
+            }
+            ```
+            """)
+            
+            gr.Markdown("""
             ## Tips for Creating n8n-Compatible JSON
             
             1. Ensure all JSON keys and values are properly quoted
@@ -687,17 +837,24 @@ def create_gradio_interface():
             """)
         
         # Set up event handlers
+        # Modify to handle fallback mode
+        def handle_message(message, chat_history, is_fallback):
+            if is_fallback:
+                return simple_fallback_response(message, chat_history)
+            else:
+                return generate_response(message, chat_history)
+                
         send_btn.click(
-            generate_response, 
-            inputs=[msg, chatbot], 
+            handle_message, 
+            inputs=[msg, chatbot, fallback_mode], 
             outputs=chatbot, 
             api_name="chat"
         )
         
         msg.submit(
-            generate_response, 
-            inputs=[msg, chatbot], 
-            outputs=chatbot, 
+            handle_message, 
+            inputs=[msg, chatbot, fallback_mode], 
+            outputs=chatbot,
             api_name=False
         )
         
@@ -707,6 +864,13 @@ def create_gradio_interface():
             api_name="load_model"
         )
         
+        use_fallback_btn.click(
+            toggle_fallback_mode,
+            inputs=[fallback_mode],
+            outputs=[fallback_mode, model_status],
+            api_name="fallback_mode"
+        )
+        
         update_prompt_btn.click(
             update_system_prompt, 
             inputs=system_prompt_input,