Upload 12 files

.gitattributes

@@ -1,35 +1,2 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,9 @@
+Apache License 2.0
+
+Copyright 2025 hmnshudhmn24
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0

README.md

@@ -1,3 +1,49 @@
----
-license: apache-2.0
----
+---
+language:
+  - en
+license: apache-2.0
+tags:
+  - code-explanation
+  - visualization
+  - mermaid
+  - codet5
+  - developer-tools
+pipeline_tag: text-generation
+library_name: transformers
+base_model: Salesforce/codet5-small
+---
+
+# code-explain-viz
+
+**Short:** `code-explain-viz` explains functions, generates step-by-step reasoning, creates a Mermaid flowchart of control flow, and suggests unit tests — combining LLM-generated explanations with deterministic AST-based visualizations.
+
+## Quick start
+
+1. Install requirements:
+```bash
+pip install -r requirements.txt
+```
+
+2. Run CLI demo:
+```bash
+python cli.py --file data_examples/example_code.py
+```
+
+3. Copy the Mermaid flowchart text printed by CLI into a Mermaid live editor (https://mermaid.live) or render with mermaid-cli to see the visual flowchart.
+
+## What you get
+- `short` one-line explanation
+- `detailed` explanation (multi-line)
+- `mermaid` flowchart text describing control flow
+- `unit_tests` template (pytest)
+
+## How it works
+- A sequence-to-sequence model (CodeT5) generates natural language explanations from code.
+- `viz_generator.py` parses the function AST and produces a reliable mermaid flowchart.
+- Combining both yields both human-friendly narrative and precise structural view.
+
+## Train / Fine-tune
+Use `train_docgen.py` with a JSONL dataset (each line: `{"code": "...", "doc": "..."}`).
+
+## License
+Apache-2.0

cli.py

@@ -0,0 +1,34 @@
+# cli.py
+import argparse
+from inference import CodeExplainViz
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--file", type=str, help="Path to python file with a function")
+    parser.add_argument("--code", type=str, help="Code string to explain")
+    parser.add_argument("--model", type=str, default="Salesforce/codet5-small", help="Model path or HF name")
+    args = parser.parse_args()
+
+    code = None
+    if args.file:
+        with open(args.file, "r", encoding="utf-8") as f:
+            code = f.read()
+    elif args.code:
+        code = args.code
+    else:
+        print("Provide --file or --code")
+        return
+
+    explainer = CodeExplainViz(model_name_or_path=args.model)
+    out = explainer.explain(code)
+    print("\n--- Short Explanation ---\n")
+    print(out["short"])
+    print("\n--- Detailed Explanation ---\n")
+    print(out["detailed"])
+    print("\n--- Mermaid Flowchart (copy into mermaid live editor) ---\n")
+    print(out["mermaid"])
+    print("\n--- Unit test template ---\n")
+    print(out["unit_tests"])
+
+if __name__ == "__main__":
+    main()

configs/model_config.json

@@ -0,0 +1,8 @@
+{
+  "base_model": "Salesforce/codet5-small",
+  "max_input_length": 512,
+  "max_target_length": 256,
+  "train_epochs": 3,
+  "train_batch_size": 8,
+  "learning_rate": 3e-05
+}

data_examples/example_code.py

@@ -0,0 +1,9 @@
+def factorial(n):
+    if not isinstance(n, int):
+        raise TypeError("n must be an integer")
+    if n < 0:
+        raise ValueError("n must be >= 0")
+    result = 1
+    for i in range(2, n+1):
+        result *= i
+    return result

inference.py

@@ -0,0 +1,44 @@
+# inference.py
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import textwrap
+from viz_generator import code_to_mermaid
+
+DEFAULT_MODEL = "Salesforce/codet5-small"
+
+class CodeExplainViz:
+    def __init__(self, model_name_or_path=DEFAULT_MODEL):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path)
+
+    def explain(self, code: str, max_length: int = 256) -> dict:
+        prompt = "explain: " + code
+        inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
+        outputs = self.model.generate(**inputs, max_length=max_length, num_beams=4, early_stopping=True)
+        text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        lines = [l.strip() for l in text.splitlines() if l.strip()]
+        short = lines[0] if lines else textwrap.shorten(text, width=120)
+        detailed = "\n".join(lines[1:]) if len(lines) > 1 else text
+        mermaid = code_to_mermaid(code)
+        unit_tests = self._make_unit_test_template(code)
+        return {"short": short, "detailed": detailed, "mermaid": mermaid, "unit_tests": unit_tests}
+
+    def _make_unit_test_template(self, code: str) -> str:
+        import re
+        m = re.search(r"def\s+([A-Za-z0-9_]+)\s*\((.*?)\):", code)
+        fn = m.group(1) if m else "function_under_test"
+        params = m.group(2) if m else ""
+        param_count = len([p for p in params.split(',') if p.strip()]) if params.strip() else 0
+        args = ", ".join(["0"] * param_count)
+        template = f"""import pytest
+
+from your_module import {fn}
+
+def test_{fn}_basic():
+    # TODO: replace with real inputs and expected outputs
+    assert {fn}({args}) == ...
+
+def test_{fn}_edge_cases():
+    # Example edge-case tests
+    with pytest.raises(Exception):
+        {fn}(...)"""
+        return template

requirements.txt

@@ -0,0 +1,7 @@
+transformers>=4.30.0
+datasets>=2.10.0
+torch>=1.12.0
+astor
+graphviz
+pytest
+black

tests/test_viz_and_explain.py

@@ -0,0 +1,13 @@
+import os
+from data_examples.example_code import factorial
+from inference import CodeExplainViz
+
+def test_explain_and_viz_runs():
+    with open("data_examples/example_code.py", "r", encoding="utf-8") as f:
+        code = f.read()
+    expl = CodeExplainViz()
+    out = expl.explain(code)
+    assert "mermaid" in out
+    assert out["mermaid"].startswith("flowchart")
+    assert isinstance(out["short"], str)
+    assert len(out["short"]) > 0

train_docgen.py

@@ -0,0 +1,43 @@
+# train_docgen.py
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, Seq2SeqTrainingArguments, Seq2SeqTrainer
+from datasets import load_dataset
+import argparse
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--data", type=str, default="data_examples/sample_dataset.jsonl", help="jsonl with {'code','doc'}")
+    p.add_argument("--output_dir", type=str, default="./code-explain-viz-model")
+    p.add_argument("--epochs", type=int, default=1)
+    return p.parse_args()
+
+def preprocess_batch(examples, tokenizer, max_src=512, max_tgt=256):
+    inputs = ["explain: " + c for c in examples["code"]]
+    model_inputs = tokenizer(inputs, truncation=True, padding="max_length", max_length=max_src)
+    labels = tokenizer(text_target=examples["doc"], truncation=True, padding="max_length", max_length=max_tgt)
+    model_inputs["labels"] = labels["input_ids"]
+    return model_inputs
+
+def main():
+    args = parse_args()
+    model_name = "Salesforce/codet5-small"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+
+    ds = load_dataset("json", data_files={"train": args.data})
+    tokenized = ds["train"].map(lambda x: preprocess_batch(x, tokenizer), batched=True, remove_columns=ds["train"].column_names)
+
+    training_args = Seq2SeqTrainingArguments(
+        output_dir=args.output_dir,
+        num_train_epochs=args.epochs,
+        per_device_train_batch_size=2,
+        save_strategy="epoch",
+        logging_steps=50
+    )
+    trainer = Seq2SeqTrainer(model=model, args=training_args, train_dataset=tokenized)
+    trainer.train()
+    trainer.save_model(args.output_dir)
+    tokenizer.save_pretrained(args.output_dir)
+    print("Saved model to", args.output_dir)
+
+if __name__ == "__main__":
+    main()

utils.py

@@ -0,0 +1,6 @@
+# utils.py
+import re
+
+def extract_first_function_name(code: str):
+    m = re.search(r"def\s+([A-Za-z0-9_]+)\s*\(", code)
+    return m.group(1) if m else None

viz_generator.py

@@ -0,0 +1,117 @@
+# viz_generator.py
+import ast
+
+class VizBuilder(ast.NodeVisitor):
+    def __init__(self):
+        self.nodes = []
+        self.edges = []
+        self.counter = 0
+
+    def new_id(self, prefix="n"):
+        self.counter += 1
+        return f"{prefix}{self.counter}"
+
+    def add_node(self, nid, label):
+        label = label.replace("\n", "\\n").replace('"', '\\"')
+        self.nodes.append((nid, label))
+
+    def add_edge(self, a, b, label=""):
+        self.edges.append((a, b, label))
+
+    def visit_FunctionDef(self, node: ast.FunctionDef):
+        start = self.new_id("start")
+        self.add_node(start, f"def {node.name}(...)")
+        prev = start
+        for stmt in node.body:
+            cur = self.visit(stmt)
+            if cur:
+                self.add_edge(prev, cur)
+                prev = cur
+        return start
+
+    def visit_Return(self, node: ast.Return):
+        nid = self.new_id("ret")
+        val = ast.unparse(node.value) if node.value else ""
+        self.add_node(nid, f"return {val}")
+        return nid
+
+    def visit_Raise(self, node: ast.Raise):
+        nid = self.new_id("raise")
+        exc = ast.unparse(node.exc) if node.exc else ""
+        self.add_node(nid, f"raise {exc}")
+        return nid
+
+    def visit_For(self, node: ast.For):
+        nid = self.new_id("for")
+        target = ast.unparse(node.target)
+        iter_ = ast.unparse(node.iter)
+        self.add_node(nid, f"for {target} in {iter_}")
+        prev = nid
+        for stmt in node.body:
+            cur = self.visit(stmt)
+            if cur:
+                self.add_edge(prev, cur)
+                prev = cur
+        return nid
+
+    def visit_While(self, node: ast.While):
+        nid = self.new_id("while")
+        cond = ast.unparse(node.test)
+        self.add_node(nid, f"while {cond}")
+        prev = nid
+        for stmt in node.body:
+            cur = self.visit(stmt)
+            if cur:
+                self.add_edge(prev, cur)
+                prev = cur
+        return nid
+
+    def visit_If(self, node: ast.If):
+        nid = self.new_id("if")
+        cond = ast.unparse(node.test)
+        self.add_node(nid, f"if {cond}")
+        for stmt in node.body:
+            cur = self.visit(stmt)
+            if cur:
+                self.add_edge(nid, cur, label="true")
+        if node.orelse:
+            for stmt in node.orelse:
+                cur = self.visit(stmt)
+                if cur:
+                    self.add_edge(nid, cur, label="false")
+        return nid
+
+    def visit_Expr(self, node: ast.Expr):
+        nid = self.new_id("expr")
+        txt = ast.unparse(node.value)
+        self.add_node(nid, txt)
+        return nid
+
+    def visit_Assign(self, node: ast.Assign):
+        nid = self.new_id("assign")
+        targets = ", ".join([ast.unparse(t) for t in node.targets])
+        val = ast.unparse(node.value)
+        self.add_node(nid, f"{targets} = {val}")
+        return nid
+
+    def generic_visit(self, node):
+        super().generic_visit(node)
+        return None
+
+def code_to_mermaid(code: str) -> str:
+    tree = ast.parse(code)
+    vb = VizBuilder()
+    root_id = None
+    for node in tree.body:
+        if isinstance(node, ast.FunctionDef):
+            root_id = vb.visit(node)
+            break
+    lines = ["flowchart TD"]
+    for nid, label in vb.nodes:
+        lines.append(f'    {nid}["{label}"]')
+    for a, b, lbl in vb.edges:
+        if lbl:
+            lines.append(f'    {a} -->|{lbl}| {b}')
+        else:
+            lines.append(f'    {a} --> {b}')
+    return "\n".join(lines)