Final local changes before pulling latest remote updates

Dockerfile

@@ -1,20 +1,27 @@
-FROM python:3.13.5-slim
+# Use a Python base image
+FROM python:3.10-slim
 
+# Set the working directory
 WORKDIR /app
 
-RUN apt-get update && apt-get install -y \
-    build-essential \
-    curl \
-    git \
-    && rm -rf /var/lib/apt/lists/*
-
-COPY requirements.txt ./
-COPY src/ ./src/
+# Expose the port Streamlit runs on
+EXPOSE 8501
 
-RUN pip3 install -r requirements.txt
+# Copy the requirements file and install dependencies
+COPY requirements.txt .
 
-EXPOSE 8501
+# Install dependencies. We use specific indexes for PyTorch to ensure compatibility 
+# and a general upgrade to avoid issues.
+RUN pip install --upgrade pip
+RUN pip install -r requirements.txt
 
-HEALTHCHECK CMD curl --fail http://localhost:8501/_stcore/health
+# Copy the rest of the application code
+# Assuming the user saved the provided code blocks into these files
+COPY main.py . 
+COPY merged.py .
+COPY app.py .
 
-ENTRYPOINT ["streamlit", "run", "src/streamlit_app.py", "--server.port=8501", "--server.address=0.0.0.0"]
+# Command to run the Streamlit application
+# We use --server.port 8501 to match the exposed port
+# and --server.address 0.0.0.0 to make it accessible outside the container
+CMD ["streamlit", "run", "app.py", "--server.port", "8501", "--server.address", "0.0.0.0"]

app.py

@@ -0,0 +1,83 @@
+import streamlit as st
+from merged import load_model_and_tokenizer, generate_code_stream
+from main import test_case   # your existing PDF generator
+
+st.write("🚀 App started")
+# ---------------- Load model ONCE ----------------
+@st.cache_resource
+def load_model():
+    return load_model_and_tokenizer()
+
+tokenizer, model = load_model()
+
+# ---------------- Session state ----------------
+for key in ["selected_language", "generated_code", "pdf_bytes", "user_input"]:
+    if key not in st.session_state:
+        st.session_state[key] = "" if key != "selected_language" else None
+
+# ---------------- UI helpers ----------------
+def select_language(lang):
+    st.session_state.selected_language = lang
+    st.session_state.generated_code = ""
+    st.session_state.pdf_bytes = ""
+    st.session_state.user_input = ""
+
+def reset():
+    for k in st.session_state:
+        st.session_state[k] = "" if k != "selected_language" else None
+
+# ---------------- UI ----------------
+st.title("Generate any code and get test case for it")
+
+if st.session_state.selected_language is None:
+    cols = st.columns(4)
+    cols[0].button("Java", on_click=select_language, args=("Java",))
+    cols[1].button("React.js", on_click=select_language, args=("React.js",))
+    cols[2].button("Python", on_click=select_language, args=("Python",))
+    cols[3].button("C++", on_click=select_language, args=("C++",))
+
+else:
+    lang = st.session_state.selected_language
+    st.subheader(f"Selected Language: {lang}")
+    st.button("Reset", on_click=reset)
+
+    st.session_state.user_input = st.text_input(
+        "Describe the task",
+        value=st.session_state.user_input
+    )
+
+    if st.button("Generate Code"):
+        if st.session_state.user_input.strip():
+            st.session_state.generated_code = ""
+
+            st.subheader("Generated Code")
+            code_placeholder = st.empty()
+
+            with st.spinner("Generating code..."):
+                for token in generate_code_stream(
+                    lang,
+                    st.session_state.user_input,
+                    tokenizer,
+                    model
+                ):
+                    st.session_state.generated_code += token
+                    code_placeholder.code(
+                        st.session_state.generated_code
+                    )
+
+            # Generate test case PDF AFTER code generation
+            st.session_state.pdf_bytes = test_case(
+                st.session_state.generated_code
+            )
+
+        else:
+            st.warning("Please enter a task")
+
+    if st.session_state.pdf_bytes:
+        st.download_button(
+            "Download Test Cases PDF",
+            st.session_state.pdf_bytes,
+            file_name="test_cases.pdf",
+            mime="application/pdf"
+        )
+

finetunned.py

@@ -0,0 +1,34 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
+import torch
+
+model_id = "Qwen/Qwen2.5-Coder-0.5B-Instruct"
+lora_path = "./Qwen2.5-Coder-0.5B-lora"
+
+tokenizer = AutoTokenizer.from_pretrained(
+    model_id,
+    trust_remote_code=True
+)
+tokenizer.pad_token = tokenizer.eos_token
+
+# 🔴 CPU ONLY — NO CUDA, NO device_map
+base_model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    trust_remote_code=True,
+    torch_dtype=torch.float32,
+    low_cpu_mem_usage=True,
+)
+
+model = PeftModel.from_pretrained(
+    base_model,
+    lora_path,
+)
+
+print("🔄 Merging LoRA (this will take time on CPU)...")
+
+merged_model = model.merge_and_unload()
+
+merged_model.save_pretrained("./Qwen2.5-Coder-0.5B-lora-merged")
+tokenizer.save_pretrained("./Qwen2.5-Coder-0.5B-lora-merged")
+
+print("✅ Merge complete")

main.py

@@ -0,0 +1,88 @@
+from langchain_core.runnables import RunnablePassthrough, RunnableMap
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompts import PromptTemplate, ChatPromptTemplate
+from langchain_huggingface import HuggingFacePipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import torch
+from fpdf import FPDF
+import re
+
+model_id = "deepseek-ai/deepseek-coder-1.3b-instruct"
+
+# Load model and tokenizer once
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    dtype=torch.float16,
+    device_map="auto",
+    offload_folder="./offload"
+)
+stop_tokens = ["<|end_of_text|>", "<|end_of_user|>"]
+
+# Wrap Transformers pipeline
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+llm = HuggingFacePipeline(
+    pipeline=pipe,
+    model_kwargs={
+        "max_new_tokens": 4096,
+        "do_sample": True,
+        "temperature": 0.2,
+        "repetition_penalty": 1.05,
+        "eos_token_id": tokenizer.eos_token_id,
+    }
+)
+
+def test_case(code):
+    test_prompt = ChatPromptTemplate.from_messages(
+        [
+            (
+                "system",
+                """You are an expert QA engineer.
+    STRICTLY follow these rules for your output:
+    - Generate EXACTLY 10 numbered test cases (1–5 functional, 6–10 edge cases).
+    - Output ONLY the numbered list.
+    - DO NOT include explanations, headers, filler text, or markdown.
+    - Each test MUST be a single, concise sentence.
+    - Begin your response immediately with '1. '""", # Slightly relaxed constraints
+            ),
+            (
+                "user",
+                "Generate test cases for the following code:\n{code}",
+            ),
+        ]
+    )
+
+    test_chain = test_prompt | llm | StrOutputParser()
+    test_cases = test_chain.invoke({"code": code})
+
+    print("\nGenerated Test Cases (Raw):\n", test_cases)
+
+    # Aggressive cleaning
+    test_cases = re.sub(r"```.*?```", "", test_cases, flags=re.DOTALL)
+    test_cases = re.sub(r"```", "", test_cases)
+    test_cases = test_cases.strip()
+
+    # --- ADD THIS CHECK ---
+    if not test_cases:
+        test_cases = "Error: Test case generation failed or returned empty content."
+    
+    print("\nGenerated Test Cases (Cleaned):\n", test_cases)
+    # -----------------------
+
+    # Encoding step remains the same for FPDF compatibility
+    safe_text = test_cases.encode("latin-1", "ignore").decode("latin-1")
+
+    # If the safe_text is still empty, FPDF will produce an empty PDF
+    # It's better to verify the content being passed to FPDF
+    
+    pdf=FPDF()
+    pdf.add_page()
+    pdf.set_font("Arial", size=12)
+    # You can set a title to ensure the PDF isn't blank
+    pdf.multi_cell(0, 10, txt="--- Generated Test Cases ---", align='C')
+    pdf.multi_cell(0, 10, txt=safe_text)
+    
+    # Use output as a bytes object without saving to disk first
+    file = pdf.output(dest='S').encode('latin-1') 
+    
+    return file

merged.py

@@ -0,0 +1,69 @@
+import torch
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    TextIteratorStreamer
+)
+from threading import Thread
+
+MODEL_PATH = "jaiwinrc7/Qwen2.5-Coder-0.5B-finetunned-merged"
+
+
+def load_model_and_tokenizer():
+    tokenizer = AutoTokenizer.from_pretrained(
+        MODEL_PATH,
+        trust_remote_code=True
+    )
+
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_PATH,
+        device_map="cpu",
+        trust_remote_code=True,
+        torch_dtype=torch.float32,   # IMPORTANT: faster on CPU
+    )
+
+    model.eval()
+    return tokenizer, model
+
+
+def build_prompt(lang, task):
+    # Keep prompt SIMPLE for speed
+    return f"""You are a coding assistant.
+            Write {lang} code for the following task:
+
+            {task}
+
+            Code:
+            """
+
+
+def generate_code_stream(lang, user_input, tokenizer, model):
+    prompt = build_prompt(lang, user_input)
+    inputs = tokenizer(prompt, return_tensors="pt")
+
+    streamer = TextIteratorStreamer(
+        tokenizer,
+        skip_prompt=True,
+        skip_special_tokens=True
+    )
+
+    generation_kwargs = dict(
+        **inputs,
+        max_new_tokens=250,     
+        do_sample=False,
+        temperature=0.0,
+        use_cache=True,
+        streamer=streamer,
+        pad_token_id=tokenizer.eos_token_id,
+    )
+
+    # Run generation in background thread
+    thread = Thread(
+        target=model.generate,
+        kwargs=generation_kwargs
+    )
+    thread.start()
+
+    # Yield tokens as they arrive
+    for token in streamer:
+        yield token