Update app.py

app.py

@@ -1,95 +1,63 @@
 import streamlit as st
 from transformers import AutoTokenizer
 from optimum.intel import OVModelForCausalLM
-import torch
-import time
 import re
 
-# --- App Config ---
-st.set_page_config(
-    page_title="Flash Jenkins Validator",
-    page_icon="⚡",
-    layout="wide" # Wide layout helps side-by-side comparison
-)
-
+# --- Setup ---
 MODEL_ID = "Qwen/Qwen2.5-Coder-0.5B-Instruct"
 
-@st.cache_resource(show_spinner="Optimizing Engine for CPU...")
-def load_model():
+st.set_page_config(page_title="Code Validator & Fixer", layout="wide")
+
+@st.cache_resource
+def load_engine():
     tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+    # export=True will convert the model to OpenVINO IR on first run
     model = OVModelForCausalLM.from_pretrained(
-        MODEL_ID,
-        export=True,
-        compile=True,
-        load_in_8bit=True, 
-        device="CPU"
+        MODEL_ID, 
+        export=True, 
+        device="CPU", 
+        load_in_8bit=True
     )
     return tokenizer, model
 
-tokenizer, model = load_model()
-
-# --- UI Layout ---
-st.title("⚡ Flash Jenkins Validator & Fixer")
-
-col1, col2 = st.columns(2)
-
-with col1:
-    st.subheader("Input")
-    jenkinsfile = st.text_area(
-        "Paste Jenkinsfile:",
-        placeholder="pipeline { ... }",
-        height=400
-    )
-    analyze_btn = st.button("Analyze & Suggest Fix", use_container_width=True)
-
-# --- Logic ---
-if analyze_btn:
-    if not jenkinsfile.strip():
-        st.error("Please paste a Jenkinsfile first.")
-    else:
-        messages = [
-            {
-                "role": "system", 
-                "content": "You are a DevOps expert. Review the Jenkinsfile for errors. ALWAYS provide a corrected version of the full Jenkinsfile inside a ```groovy code block at the end of your response."
-            },
-            {"role": "user", "content": f"Fix this Jenkinsfile:\n\n{jenkinsfile}"}
-        ]
+tokenizer, model = load_engine()
+
+# --- UI ---
+st.title("🏗️ Smart Code Validator")
+st.write("Analyzes and fixes your code using Qwen2.5-Coder (CPU-Optimized)")
+
+col_input, col_output = st.columns(2)
+
+with col_input:
+    user_code = st.text_area("Paste your code here:", height=400, placeholder="pipeline { ... }")
+    run_btn = st.button("Validate & Fix", use_container_width=True)
+
+if run_btn and user_code:
+    # 1. Prepare Request
+    messages = [
+        {"role": "system", "content": "You are a senior developer. First, explain errors briefly. Second, provide the full corrected code inside a ```groovy block."},
+        {"role": "user", "content": f"Review and fix this code:\n{user_code}"}
+    ]
+    
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer(prompt, return_tensors="pt")
+
+    # 2. Generate
+    with st.spinner("Analyzing on CPU..."):
+        tokens = model.generate(**inputs, max_new_tokens=1024, do_sample=False)
+        response = tokenizer.decode(tokens[0][len(inputs[0]):], skip_special_tokens=True)
+    
+    # 3. Extract Fixed Code
+    code_match = re.search(r"```(?:\w+)?\n([\s\S]*?)```", response)
+    fixed_code = code_match.group(1) if code_match else None
+    explanation = re.sub(r"```[\s\S]*?```", "", response).strip()
+
+    with col_output:
+        st.subheader("Analysis")
+        st.info(explanation if explanation else "Code looks good! Here is the optimized version:")
         
-        input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        inputs = tokenizer(input_text, return_tensors="pt")
-
-        with st.status("Generating Fix...", expanded=True) as status:
-            start_time = time.time()
-            
-            output_tokens = model.generate(
-                **inputs,
-                max_new_tokens=1024, # Increased to ensure full code block fits
-                do_sample=False,
-                pad_token_id=tokenizer.eos_token_id
-            )
-            
-            new_tokens = output_tokens[0][len(inputs["input_ids"][0]):]
-            response = tokenizer.decode(new_tokens, skip_special_tokens=True)
-            
-            # --- Feature: Extraction Logic ---
-            # Search for the code block in the AI response
-            code_match = re.search(r"```groovy\n([\s\S]*?)```", response)
-            fixed_code = code_match.group(1) if code_match else None
-            # Clean up explanation by removing the code block from the markdown view
-            explanation = re.sub(r"```groovy[\s\S]*?```", "*(See corrected code in the Fix panel)*", response)
-
-            with col2:
-                st.subheader("Analysis & Fixed Code")
-                st.markdown(explanation)
-                
-                if fixed_code:
-                    st.success("✅ One-Click Fix Ready:")
-                    st.code(fixed_code, language="groovy")
-                    st.button("Apply Fix to Input", on_click=lambda: st.session_state.update({"input_area": fixed_code}), key="apply_fix")
-                else:
-                    st.warning("Could not extract a clean code block. Check the explanation.")
-
-            status.update(label=f"Done in {time.time() - start_time:.2f}s", state="complete")
+        if fixed_code:
+            st.subheader("Corrected Code")
+            st.code(fixed_code)
+            st.download_button("Download Fixed File", fixed_code, file_name="fixed_code.groovy")
 
-st.divider()
-st.caption(f"Engine: {MODEL_ID} | CPU Optimized")