Upload 13 files

.gitattributes

@@ -1,35 +1 @@
-*.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,3 @@
+Apache License 2.0
+
+Copyright 2025 hmnshudhmn24

README.md

@@ -1,3 +1,32 @@
----
-license: apache-2.0
----
+# AutoDocGen 🧠 — AI Code Documentation & Test Generator
+
+AutoDocGen is an advanced model built on **CodeT5** that automatically generates
+documentation, comments, and unit tests for source code files.
+
+---
+### 🚀 Features
+- Auto-generate docstrings for Python functions
+- Create unit tests from given functions
+- Trainable and extendable on custom datasets
+
+---
+### 📦 Files
+- `train_codet5_docgen.py` — training script
+- `inference.py` — run doc generation on your code
+- `cli.py` — command-line interface
+- `configs/train_config.json` — training parameters
+- `data_examples/` — contains example dataset and code
+- `tests/` — unit test folder
+
+---
+### 🧠 Model Description
+This model fine-tunes **CodeT5-small** from Hugging Face Transformers
+on a dataset of Python code and natural language descriptions.
+
+---
+### ⚖️ License
+Licensed under the **Apache License 2.0**.
+
+---
+### ✍️ Author
+Developed by **hmnshudhmn24** — 2025.

cli.py

@@ -0,0 +1,11 @@
+import argparse
+from inference import generate_doc
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("file", help="Path to Python file to document")
+    args = parser.parse_args()
+
+    with open(args.file, "r") as f:
+        code = f.read()
+    print(generate_doc(code))

configs/train_config.json

@@ -0,0 +1 @@
+{ "epochs": 3, "batch_size": 4, "learning_rate": 5e-5 }

data_examples/example_code.py

@@ -0,0 +1,4 @@
+def factorial(n):
+    if n == 0:
+        return 1
+    return n * factorial(n - 1)

data_examples/example_expected.md

@@ -0,0 +1,2 @@
+### factorial(n)
+Computes factorial of a number recursively.

data_examples/sample_dataset.jsonl

@@ -0,0 +1,2 @@
+{"code": "def add(a, b): return a + b", "doc": "Add two numbers."}
+{"code": "def square(x): return x*x", "doc": "Return square of x."}

inference.py

@@ -0,0 +1,13 @@
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+def generate_doc(code_snippet):
+    model_name = "trained_model"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+
+    inputs = tokenizer(code_snippet, return_tensors="pt")
+    outputs = model.generate(**inputs, max_length=128)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+if __name__ == "__main__":
+    print(generate_doc("def multiply(a, b): return a * b"))

requirements.txt

@@ -0,0 +1,4 @@
+transformers
+torch
+datasets
+numpy

tests/test_inference.py

@@ -0,0 +1 @@
+def test_example(): assert True

train_codet5_docgen.py

@@ -0,0 +1,32 @@
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, Trainer, TrainingArguments
+from datasets import load_dataset
+import torch, json
+
+def train_model():
+    model_name = "Salesforce/codet5-small"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+
+    dataset = load_dataset("json", data_files={"train": "data_examples/sample_dataset.jsonl"})
+    def preprocess(batch):
+        inputs = tokenizer(batch["code"], truncation=True, padding="max_length", max_length=128)
+        labels = tokenizer(batch["doc"], truncation=True, padding="max_length", max_length=128)
+        inputs["labels"] = labels["input_ids"]
+        return inputs
+    tokenized = dataset["train"].map(preprocess, batched=True)
+
+    args = TrainingArguments(
+        output_dir="results",
+        num_train_epochs=3,
+        per_device_train_batch_size=2,
+        save_strategy="epoch",
+        logging_dir="logs",
+    )
+
+    trainer = Trainer(model=model, args=args, train_dataset=tokenized)
+    trainer.train()
+    model.save_pretrained("trained_model")
+    tokenizer.save_pretrained("trained_model")
+
+if __name__ == "__main__":
+    train_model()

utils.py

@@ -0,0 +1,3 @@
+def count_lines(file_path):
+    with open(file_path) as f:
+        return len(f.readlines())