Update app.py

app.py

@@ -1,322 +1,59 @@
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import json
-import gradio as gr
-from typing import Optional, Dict, Any
 
-class QwenN8NGenerator:
-    def __init__(self, model_name: str = "npv2k1/Qwen2.5-7B-n8n"):
-        """
-        Initialize the Qwen2.5-7B-n8n model for generating n8n workflows.
-        
-        Args:
-            model_name: HuggingFace model identifier
-        """
-        self.model_name = model_name
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        
-        print(f"Loading model on {self.device}...")
-        
-        try:
-            # Try loading tokenizer from the model
-            self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-        except (OSError, ValueError) as e:
-            print(f"Failed to load tokenizer from {model_name}, using base Qwen2.5 tokenizer...")
-            # Fallback to base Qwen2.5 tokenizer
-            self.tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct", trust_remote_code=True)
-        
-        try:
-            self.model = AutoModelForCausalLM.from_pretrained(
-                model_name,
-                torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
-                device_map="auto" if self.device == "cuda" else None,
-                trust_remote_code=True
-            )
-        except Exception as e:
-            print(f"Error loading model: {e}")
-            raise
-        
-        # Set padding token if not exists
-        if self.tokenizer.pad_token is None:
-            self.tokenizer.pad_token = self.tokenizer.eos_token
-            
-        print("Model loaded successfully!")
-    
-    def generate_n8n_workflow(
-        self, 
-        prompt: str, 
-        max_length: int = 2048,
-        temperature: float = 0.7,
-        top_p: float = 0.9,
-        do_sample: bool = True
-    ) -> str:
-        """
-        Generate n8n workflow JSON from a natural language prompt.
-        
-        Args:
-            prompt: Natural language description of the workflow
-            max_length: Maximum tokens to generate
-            temperature: Sampling temperature
-            top_p: Top-p sampling parameter
-            do_sample: Whether to use sampling
-            
-        Returns:
-            Generated n8n workflow JSON as string
-        """
-        # Format prompt for the model
-        formatted_prompt = f"Create an n8n workflow for: {prompt}\n\nWorkflow JSON:"
-        
-        # Tokenize input
-        inputs = self.tokenizer(
-            formatted_prompt, 
-            return_tensors="pt",
-            truncation=True,
-            max_length=512
-        ).to(self.device)
-        
-        # Generate response
-        with torch.no_grad():
-            outputs = self.model.generate(
-                **inputs,
-                max_length=max_length,
-                temperature=temperature,
-                top_p=top_p,
-                do_sample=do_sample,
-                pad_token_id=self.tokenizer.eos_token_id,
-                eos_token_id=self.tokenizer.eos_token_id,
-                repetition_penalty=1.1
-            )
-        
-        # Decode and clean output
-        generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-        
-        # Extract JSON part (assuming it starts after the prompt)
-        json_start = generated_text.find('{')
-        if json_start != -1:
-            workflow_json = generated_text[json_start:]
-            
-            # Try to validate JSON
-            try:
-                parsed = json.loads(workflow_json)
-                return json.dumps(parsed, indent=2)
-            except json.JSONDecodeError:
-                # Return raw output if JSON parsing fails
-                return workflow_json
-        
-        return generated_text
-    
-    def enhance_existing_workflow(
-        self, 
-        existing_workflow: str, 
-        enhancement_request: str,
-        max_length: int = 2048,
-        temperature: float = 0.7
-    ) -> str:
-        """
-        Enhance an existing n8n workflow based on a request.
-        
-        Args:
-            existing_workflow: Existing n8n workflow JSON
-            enhancement_request: Description of desired enhancements
-            max_length: Maximum tokens to generate
-            temperature: Sampling temperature
-            
-        Returns:
-            Enhanced n8n workflow JSON as string
-        """
-        prompt = f"""Enhance this n8n workflow: {enhancement_request}
-
-Current workflow:
-{existing_workflow}
-
-Enhanced workflow JSON:"""
-        
-        inputs = self.tokenizer(
-            prompt, 
-            return_tensors="pt",
-            truncation=True,
-            max_length=1024
-        ).to(self.device)
-        
-        with torch.no_grad():
-            outputs = self.model.generate(
-                **inputs,
-                max_length=max_length,
-                temperature=temperature,
-                top_p=0.9,
-                do_sample=True,
-                pad_token_id=self.tokenizer.eos_token_id,
-                repetition_penalty=1.1
-            )
-        
-        generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-        
-        # Extract enhanced workflow
-        json_start = generated_text.rfind('{')
-        if json_start != -1:
-            enhanced_workflow = generated_text[json_start:]
-            try:
-                parsed = json.loads(enhanced_workflow)
-                return json.dumps(parsed, indent=2)
-            except json.JSONDecodeError:
-                return enhanced_workflow
-        
-        return generated_text
+def load_qwen_n8n_model():
+    """Load the Qwen2.5-7B-n8n model with fallback tokenizer"""
+    model_name = "npv2k1/Qwen2.5-7B-n8n"
+    
+    # Load tokenizer with fallback
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+    except:
+        print("Using base Qwen tokenizer...")
+        tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct", trust_remote_code=True)
+    
+    # Load model
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        device_map="auto" if torch.cuda.is_available() else None,
+        trust_remote_code=True
+    )
+    
+    return tokenizer, model
 
-# Usage example
-def main():
-    # Initialize the generator
-    generator = QwenN8NGenerator()
-    
-    # Example 1: Generate new workflow
-    prompt = "Send a daily email with weather forecast to my team using OpenWeatherMap API"
-    workflow = generator.generate_n8n_workflow(prompt)
-    print("Generated Workflow:")
-    print(workflow)
-    print("\n" + "="*50 + "\n")
-    
-    # Example 2: Enhance existing workflow
-    simple_workflow = """{
-        "nodes": [
-            {
-                "id": "1",
-                "name": "Start",
-                "type": "n8n-nodes-base.start"
-            }
-        ],
-        "connections": {}
-    }"""
-    
-    enhancement = "Add error handling and logging"
-    enhanced = generator.enhance_existing_workflow(simple_workflow, enhancement)
-    print("Enhanced Workflow:")
-    print(enhanced)
-
-# Gradio Interface
-def create_gradio_interface():
-    generator = QwenN8NGenerator()
-    
-    def generate_workflow_ui(prompt, max_length, temperature):
-        try:
-            workflow = generator.generate_n8n_workflow(
-                prompt, 
-                max_length=max_length, 
-                temperature=temperature
-            )
-            return workflow
-        except Exception as e:
-            return f"Error: {str(e)}"
-    
-    def enhance_workflow_ui(existing_workflow, enhancement_request, max_length, temperature):
-        try:
-            enhanced = generator.enhance_existing_workflow(
-                existing_workflow,
-                enhancement_request,
-                max_length=max_length,
-                temperature=temperature
-            )
-            return enhanced
-        except Exception as e:
-            return f"Error: {str(e)}"
-    
-    # Create Gradio interface
-    with gr.Blocks(title="Qwen2.5-7B n8n Workflow Generator") as demo:
-        gr.Markdown("# n8n Workflow Generator using Qwen2.5-7B")
-        
-        with gr.Tab("Generate New Workflow"):
-            with gr.Row():
-                with gr.Column():
-                    prompt_input = gr.Textbox(
-                        label="Workflow Description",
-                        placeholder="Describe the workflow you want to create...",
-                        lines=3
-                    )
-                    max_length_slider = gr.Slider(
-                        minimum=512,
-                        maximum=4096,
-                        value=2048,
-                        label="Max Length"
-                    )
-                    temperature_slider = gr.Slider(
-                        minimum=0.1,
-                        maximum=1.0,
-                        value=0.7,
-                        label="Temperature"
-                    )
-                    generate_btn = gr.Button("Generate Workflow", variant="primary")
-                
-                with gr.Column():
-                    workflow_output = gr.Code(
-                        label="Generated n8n Workflow",
-                        language="json",
-                        lines=20
-                    )
-            
-            generate_btn.click(
-                generate_workflow_ui,
-                inputs=[prompt_input, max_length_slider, temperature_slider],
-                outputs=workflow_output
-            )
-        
-        with gr.Tab("Enhance Existing Workflow"):
-            with gr.Row():
-                with gr.Column():
-                    existing_workflow_input = gr.Code(
-                        label="Existing Workflow JSON",
-                        language="json",
-                        lines=10
-                    )
-                    enhancement_input = gr.Textbox(
-                        label="Enhancement Request",
-                        placeholder="Describe how you want to enhance the workflow...",
-                        lines=3
-                    )
-                    enhance_max_length = gr.Slider(
-                        minimum=512,
-                        maximum=4096,
-                        value=2048,
-                        label="Max Length"
-                    )
-                    enhance_temperature = gr.Slider(
-                        minimum=0.1,
-                        maximum=1.0,
-                        value=0.7,
-                        label="Temperature"
-                    )
-                    enhance_btn = gr.Button("Enhance Workflow", variant="primary")
-                
-                with gr.Column():
-                    enhanced_output = gr.Code(
-                        label="Enhanced n8n Workflow",
-                        language="json",
-                        lines=20
-                    )
-            
-            enhance_btn.click(
-                enhance_workflow_ui,
-                inputs=[existing_workflow_input, enhancement_input, enhance_max_length, enhance_temperature],
-                outputs=enhanced_output
-            )
-        
-        # Examples
-        gr.Markdown("## Example Prompts")
-        gr.Examples(
-            examples=[
-                ["Create a workflow that monitors a GitHub repository for new issues and sends Slack notifications"],
-                ["Build an automated lead scoring system that processes form submissions and updates CRM"],
-                ["Design a workflow for social media automation that posts content across multiple platforms"],
-                ["Create a data pipeline that fetches API data, processes it, and stores in database"],
-                ["Build a customer support workflow that categorizes tickets and assigns to appropriate teams"]
-            ],
-            inputs=prompt_input
+def generate_n8n_workflow(tokenizer, model, prompt, max_length=1024):
+    """Generate n8n workflow from prompt"""
+    formatted_prompt = f"Create an n8n workflow: {prompt}\n\nJSON:"
+    
+    inputs = tokenizer(formatted_prompt, return_tensors="pt", truncation=True)
+    
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_length=max_length,
+            temperature=0.7,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
         )
     
-    return demo
+    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # Extract JSON
+    json_start = result.find('{')
+    if json_start != -1:
+        return result[json_start:]
+    return result
 
+# Usage
 if __name__ == "__main__":
-    # Run the main example
-    main()
+    tokenizer, model = load_qwen_n8n_model()
+    
+    workflow = generate_n8n_workflow(
+        tokenizer, 
+        model, 
+        "Send email when new GitHub issue is created"
+    )
     
-    # Uncomment to launch Gradio interface
-    # demo = create_gradio_interface()
-    # demo.launch(share=True)
+    print(workflow)