Update app.py

app.py

@@ -17,26 +17,6 @@ import re
 MODEL_REPO = "Nishan30/n8n-workflow-generator"  # Update with your HF repo
 BASE_MODEL = "Qwen/Qwen2.5-Coder-1.5B-Instruct"
 
-SYSTEM_PROMPT = """You are an expert n8n workflow generator. Given a user's request, you MUST generate clean, functional TypeScript code that EXACTLY matches their specific requirements using the @n8n-generator/core DSL.
-
-CRITICAL: Generate code based ONLY on the user's request. Do NOT use example workflows. Create the workflow the user asks for.
-
-Your output should:
-- Only contain the code, no explanations
-- Use the Workflow class from @n8n-generator/core
-- Use workflow.add() to create nodes with appropriate parameters
-- Use .to() or workflow.connect() for connections
-- Match the user's specific requirements exactly
-- Be ready to compile directly to n8n JSON
-
-Format:
-```typescript
-const workflow = new Workflow('Descriptive Name');
-const node1 = workflow.add('n8n-nodes-base.nodetype', { param: 'value' });
-const node2 = workflow.add('n8n-nodes-base.nodetype', { param: 'value' });
-node1.to(node2);
-```"""
-
 # ==============================================================================
 # MODEL LOADING
 # ==============================================================================
@@ -120,41 +100,46 @@ def generate_workflow(prompt, temperature=0.5, max_tokens=1024):
     if not prompt.strip():
         return "Please enter a workflow description.", None, None
 
-    # Format messages
-    messages = [
-        {"role": "system", "content": SYSTEM_PROMPT},
-        {"role": "user", "content": prompt}
-    ]
-
-    # Apply chat template
-    text = tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True
-    )
+    # IMPORTANT: Use the exact format the model was trained with
+    formatted_prompt = f"""### System:
+You are an expert n8n workflow generator. Given a user's request, you generate clean, functional TypeScript code using the @n8n-generator/core DSL.
+
+Your output should:
+- Only contain the code, no explanations
+- Use the Workflow class from @n8n-generator/core
+- Use workflow.add() to create nodes
+- Use .to() or workflow.connect() for connections
+- Be ready to compile directly to n8n JSON
+
+### Instruction:
+{prompt}
+
+### Response:
+"""
 
     # Debug: Print formatted prompt (first 500 chars)
     print(f"\n{'='*60}")
     print(f"User Prompt: {prompt}")
-    print(f"Formatted Input (truncated):\n{text[:500]}...")
+    print(f"Formatted Input (truncated):\n{formatted_prompt[:500]}...")
     print(f"{'='*60}\n")
 
     # Tokenize
-    inputs = tokenizer(text, return_tensors="pt").to(model.device)
+    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
     input_length = inputs.input_ids.shape[1]
     print(f"Input tokens: {input_length}, Max new tokens: {max_tokens}")
 
-    # Generate with better sampling parameters
+    # Generate with parameters matching training
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
             max_new_tokens=max_tokens,
-            temperature=max(temperature, 0.1),  # Ensure minimum temperature
-            do_sample=True,  # Always sample for variety
-            top_p=0.9,
-            top_k=50,  # Add top-k sampling
-            repetition_penalty=1.15,  # Increase to reduce repetition
-            no_repeat_ngram_size=3  # Prevent repeating 3-grams
+            temperature=max(temperature, 0.1),
+            do_sample=True,
+            top_p=0.95,
+            top_k=50,
+            repetition_penalty=1.1,
+            eos_token_id=tokenizer.eos_token_id,
+            pad_token_id=tokenizer.pad_token_id,
         )
 
     # Decode
@@ -162,10 +147,10 @@ def generate_workflow(prompt, temperature=0.5, max_tokens=1024):
 
     # Debug: Print generated text
     print(f"Generated text length: {len(generated_text)} chars")
-    print(f"Generated text (first 300 chars):\n{generated_text[:300]}...\n")
+    print(f"Generated text (first 500 chars):\n{generated_text[:500]}...\n")
 
-    # Extract code from response
-    code = extract_code(generated_text)
+    # Extract code from response (handle ### Response: format)
+    code = extract_code_from_instruction_format(generated_text)
 
     # Convert to n8n JSON
     n8n_json = convert_to_n8n_json(code)
@@ -175,25 +160,33 @@ def generate_workflow(prompt, temperature=0.5, max_tokens=1024):
 
     return code, json.dumps(n8n_json, indent=2), visualization
 
-def extract_code(text):
-    """Extract TypeScript code from generated text"""
+def extract_code_from_instruction_format(text):
+    """Extract TypeScript code from ### Response: format"""
 
-    # Try to find code block
-    code_match = re.search(r'```(?:typescript|ts)?\n(.*?)```', text, re.DOTALL)
+    # Split by ### Response: and get the part after it
+    try:
+        response_part = text.split("### Response:")[-1].strip()
+    except:
+        response_part = text
+
+    # Remove any subsequent ### markers (like ### Instruction:, ### System:)
+    for stop_marker in ["### Instruction:", "### System:", "\n\n\n\n"]:
+        if stop_marker in response_part:
+            response_part = response_part.split(stop_marker)[0].strip()
+
+    # Try to extract code from markdown blocks
+    code_match = re.search(r'```(?:typescript|ts)?\n(.*?)```', response_part, re.DOTALL)
     if code_match:
         return code_match.group(1).strip()
 
-    # If no code block, look for code after assistant response
-    if "assistant" in text.lower():
-        parts = text.split("assistant", 1)
-        if len(parts) > 1:
-            # Remove any markdown code blocks
-            code = parts[1].strip()
-            code = re.sub(r'```(?:typescript|ts)?\n', '', code)
-            code = re.sub(r'```', '', code)
-            return code.strip()
+    # Remove markdown code block markers if present
+    response_part = re.sub(r'```(?:typescript|ts)?', '', response_part)
 
-    return text.strip()
+    return response_part.strip()
+
+def extract_code(text):
+    """Legacy extraction function - kept for compatibility"""
+    return extract_code_from_instruction_format(text)
 
 # ==============================================================================
 # N8N JSON CONVERSION
@@ -495,4 +488,4 @@ if __name__ == "__main__":
         server_name="0.0.0.0",
         server_port=7860,
         share=False
-    )
+    )