Upload folder using huggingface_hub

README.md

@@ -1,12 +1,25 @@
 ---
-title: Codewraith
-emoji: 🏃
+title: CodeWraith
+emoji: 🔮
 colorFrom: purple
-colorTo: purple
+colorTo: indigo
 sdk: gradio
-sdk_version: 6.11.0
+sdk_version: 4.44.1
 app_file: app.py
 pinned: false
+license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# CodeWraith - Module-to-Spec Transformer
+
+Generate technical specifications from Python source code using a fine-tuned LLM.
+
+## How it works
+
+1. Paste Python source code in the left panel
+2. Adjust sampling parameters (temperature, top_p, max tokens)
+3. Toggle RAG to include similar examples as context
+4. Click **Generate Specification**
+
+The model is a LoRA-fine-tuned Llama that was trained on 200+ Python module / specification pairs
+generated by a teacher model and verified with AST-based structural validation.

app.py

@@ -0,0 +1,254 @@
+"""CodeWraith HuggingFace Spaces entry point.
+
+Downloads the LoRA adapter from HF Hub and serves the Gradio interface.
+Set HF_REPO_ID environment variable to point to your uploaded adapter.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+import gradio as gr
+
+# --- Config ---
+
+HF_REPO_ID = os.environ.get("HF_REPO_ID", "slenk/codewraith-lora-3b")
+MODEL_KEY = os.environ.get("MODEL_KEY", "3b")
+ADAPTER_DIR = "./adapter"
+
+MODELS = {
+    "3b": "unsloth/Llama-3.2-3B-Instruct",
+    "8b": "unsloth/Llama-3.1-8B-Instruct",
+}
+
+SYSTEM_MESSAGE = (
+    "You are CodeWraith, a technical specification generator. "
+    "Given Python source code, produce a structured Markdown specification "
+    "that accurately captures all functions, classes, parameters, return types, "
+    "dependencies, and error handling patterns."
+)
+
+EXAMPLE_CODE = '''\
+def fibonacci(n: int) -> list[int]:
+    """Generate the first n Fibonacci numbers."""
+    if n <= 0:
+        return []
+    sequence = [0, 1]
+    while len(sequence) < n:
+        sequence.append(sequence[-1] + sequence[-2])
+    return sequence[:n]
+'''
+
+# --- Global state ---
+
+_model = None
+_tokenizer = None
+_retriever = None
+
+
+# --- Model loading ---
+
+
+def download_adapter():
+    """Download the LoRA adapter from HF Hub if not already cached."""
+    if Path(ADAPTER_DIR).exists() and any(Path(ADAPTER_DIR).iterdir()):
+        print(f"Adapter already cached at {ADAPTER_DIR}")
+        return
+
+    from huggingface_hub import snapshot_download
+
+    print(f"Downloading adapter from {HF_REPO_ID}...")
+    snapshot_download(repo_id=HF_REPO_ID, local_dir=ADAPTER_DIR)
+    print("Download complete.")
+
+
+def load_model() -> tuple[Any, Any]:
+    """Load the base model with LoRA adapter."""
+    global _model, _tokenizer  # noqa: PLW0603
+
+    if _model is not None:
+        return _model, _tokenizer
+
+    download_adapter()
+
+    from peft import PeftModel
+    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+
+    model_name = MODELS[MODEL_KEY]
+    print(f"Loading {model_name}...")
+
+    bnb_config = BitsAndBytesConfig(load_in_4bit=True)
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        quantization_config=bnb_config,
+        device_map="auto",
+    )
+    model = PeftModel.from_pretrained(model, ADAPTER_DIR)
+    model.eval()
+
+    _model, _tokenizer = model, tokenizer
+    return model, tokenizer
+
+
+# --- RAG ---
+
+
+def init_retriever():
+    """Initialize retriever if training data is bundled."""
+    global _retriever  # noqa: PLW0603
+
+    if _retriever is not None:
+        return _retriever
+
+    index_path = Path("chromadb")
+    data_path = Path("training_pairs_clean.jsonl")
+
+    if not index_path.exists() and data_path.exists():
+        # Build index from bundled data
+        try:
+            import chromadb
+            from chromadb.utils import embedding_functions
+
+            client = chromadb.PersistentClient(path=str(index_path))
+            ef = embedding_functions.SentenceTransformerEmbeddingFunction(
+                model_name="all-MiniLM-L6-v2"
+            )
+            collection = client.get_or_create_collection(
+                name="codewraith_specs", embedding_function=ef
+            )
+
+            if collection.count() == 0:
+                pairs = []
+                with data_path.open() as f:
+                    for line in f:
+                        if line.strip():
+                            pairs.append(json.loads(line))
+
+                for i in range(0, len(pairs), 50):
+                    batch = pairs[i : i + 50]
+                    collection.add(
+                        ids=[f"pair_{i + j}" for j in range(len(batch))],
+                        documents=[p["input"] for p in batch],
+                        metadatas=[{"spec": p["output"]} for p in batch],
+                    )
+
+            _retriever = (client, collection, ef)
+        except Exception as e:
+            print(f"RAG init failed: {e}")
+
+    return _retriever
+
+
+def retrieve_context(source_code: str, n_results: int = 3) -> str:
+    """Retrieve similar examples as context."""
+    ret = init_retriever()
+    if ret is None:
+        return ""
+
+    _, collection, _ = ret
+    results = collection.query(query_texts=[source_code], n_results=n_results)
+
+    parts = ["Here are examples of Python code and their specifications:\n"]
+    for i, (doc, meta) in enumerate(zip(results["documents"][0], results["metadatas"][0]), 1):
+        parts.append(
+            f"\n--- Example {i} ---\n"
+            f"Code:\n```python\n{doc[:1500]}\n```\n"
+            f"Specification:\n{meta['spec'][:1500]}\n"
+        )
+    parts.append("\nNow generate a specification for the following code:\n")
+    return "".join(parts)
+
+
+# --- Inference ---
+
+
+def generate_spec(
+    source_code: str,
+    temperature: float = 0.7,
+    top_p: float = 0.9,
+    max_tokens: int = 2048,
+    use_rag: bool = True,
+) -> str:
+    """Generate a technical specification."""
+    if not source_code.strip():
+        return "*Please paste some Python source code.*"
+
+    model, tokenizer = load_model()
+
+    user_content = source_code
+    if use_rag:
+        context = retrieve_context(source_code)
+        if context:
+            user_content = context + source_code
+
+    messages = [
+        {"role": "system", "content": SYSTEM_MESSAGE},
+        {"role": "user", "content": user_content},
+    ]
+
+    inputs = tokenizer.apply_chat_template(
+        messages, tokenize=True, add_generation_prompt=True, return_tensors="pt"
+    ).to(model.device)
+
+    outputs = model.generate(
+        input_ids=inputs,
+        max_new_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        do_sample=True,
+    )
+
+    generated = outputs[0][inputs.shape[-1] :]
+    return tokenizer.decode(generated, skip_special_tokens=True)
+
+
+# --- Gradio UI ---
+
+
+def create_app():
+    with gr.Blocks(
+        title="CodeWraith - Module-to-Spec Transformer",
+        theme=gr.themes.Soft(),
+    ) as app:
+        gr.Markdown(
+            "# CodeWraith\n"
+            "Generate technical specifications from Python source code.\n\n"
+            "Paste your Python code on the left, adjust sampling parameters, "
+            "and click **Generate Specification**."
+        )
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                code_input = gr.Code(
+                    language="python",
+                    label="Python Source Code",
+                    value=EXAMPLE_CODE,
+                    lines=20,
+                )
+                with gr.Row():
+                    temperature = gr.Slider(0.0, 2.0, value=0.7, step=0.1, label="Temperature")
+                    top_p = gr.Slider(0.0, 1.0, value=0.9, step=0.05, label="Top-p")
+                    max_tokens = gr.Slider(256, 4096, value=2048, step=256, label="Max Tokens")
+                use_rag = gr.Checkbox(value=True, label="Use RAG (retrieve similar examples)")
+                generate_btn = gr.Button("Generate Specification", variant="primary")
+
+            with gr.Column(scale=1):
+                spec_output = gr.Markdown(label="Generated Specification")
+
+        generate_btn.click(
+            fn=generate_spec,
+            inputs=[code_input, temperature, top_p, max_tokens, use_rag],
+            outputs=spec_output,
+        )
+
+    return app
+
+
+if __name__ == "__main__":
+    app = create_app()
+    app.launch()

requirements.txt

@@ -0,0 +1,8 @@
+transformers>=4.40.0
+torch>=2.1.0
+accelerate>=0.27.0
+peft>=0.10.0
+gradio>=4.0.0
+chromadb>=0.5.0
+sentence-transformers>=3.0.0
+bitsandbytes>=0.43.0