1 changed files with 565 additions and 0 deletions
@ -0,0 +1,565 @@
|
||||
{ |
||||
"cells": [ |
||||
{ |
||||
"cell_type": "markdown", |
||||
"id": "ff022957-2e81-4ea9-84d3-e52d5753e133", |
||||
"metadata": {}, |
||||
"source": [ |
||||
"### Comment and Unit Test Generater \n", |
||||
"\n", |
||||
"The requirement: \n", |
||||
"* use an LLM to generate docstring and comments for Python code\n", |
||||
"* use an LLM to generate unit test\n", |
||||
"\n", |
||||
"This is my week 4 day 5 project." |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "ea1841f6-4afc-4d29-ace8-5ca5a3915c8c", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"# imports\n", |
||||
"\n", |
||||
"import os\n", |
||||
"import io\n", |
||||
"import sys\n", |
||||
"import json\n", |
||||
"import requests\n", |
||||
"from dotenv import load_dotenv\n", |
||||
"from openai import OpenAI\n", |
||||
"import google.generativeai\n", |
||||
"import anthropic\n", |
||||
"from IPython.display import Markdown, display, update_display\n", |
||||
"import gradio as gr\n", |
||||
"import subprocess\n", |
||||
"from huggingface_hub import login, InferenceClient\n", |
||||
"from transformers import AutoTokenizer" |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "11957fd3-6c61-4496-aef1-8223cb9ec4ce", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"# environment\n", |
||||
"\n", |
||||
"load_dotenv()\n", |
||||
"os.environ['OPENAI_API_KEY'] = os.getenv('OPENAI_API_KEY', 'your-key-if-not-using-env')\n", |
||||
"os.environ['ANTHROPIC_API_KEY'] = os.getenv('ANTHROPIC_API_KEY', 'your-key-if-not-using-env')\n", |
||||
"os.environ['HF_TOKEN'] = os.getenv('HF_TOKEN', 'your-key-if-not-using-env')" |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "ee7b08fd-e678-4234-895e-4e3a925e60f0", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"# initialize\n", |
||||
"\n", |
||||
"openai = OpenAI()\n", |
||||
"claude = anthropic.Anthropic()\n", |
||||
"OPENAI_MODEL = \"gpt-4o\"\n", |
||||
"CLAUDE_MODEL = \"claude-3-5-sonnet-20240620\"" |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "c8023255-9c98-4fbc-92e4-c553bed3b605", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"hf_token = os.environ['HF_TOKEN']\n", |
||||
"login(hf_token, add_to_git_credential=True)" |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "f8ce3f5e-74c4-4d35-bfbc-91c5be85e094", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"code_qwen = \"Qwen/CodeQwen1.5-7B-Chat\"\n", |
||||
"CODE_QWEN_URL = \"https://g39mbjooiiwkbgyz.us-east-1.aws.endpoints.huggingface.cloud\"" |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "1bbc66b6-52ae-465e-a368-edc8f097fe9d", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"def system_prompt_for_comment():\n", |
||||
" system=\"\"\"\n", |
||||
" You are a Python documentation expert. When writing documentation:\n", |
||||
" - Follow PEP 257 and Google docstring style guidelines\n", |
||||
" - Write clear, concise explanations\n", |
||||
" - Include practical examples\n", |
||||
" - Highlight edge cases and limitations\n", |
||||
" - Use type hints in docstrings\n", |
||||
" - Add inline comments only for complex logic\n", |
||||
" - Never skip documenting parameters or return values\n", |
||||
" - Validate that all documentation is accurate and complete\n", |
||||
" \"\"\"\n", |
||||
" return system" |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "b089f87b-53ae-40ad-8d06-b9924bb998a0", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"def system_prompt_for_unit_test():\n", |
||||
" system=\"\"\"\n", |
||||
" You are an expert Python testing engineer who specializes in creating comprehensive unit tests. Follow these principles:\n", |
||||
" - Use pytest as the testing framework\n", |
||||
" - Follow the Arrange-Act-Assert pattern\n", |
||||
" - Test both valid and invalid inputs\n", |
||||
" - Include edge cases and boundary conditions\n", |
||||
" - Write descriptive test names that explain the scenario being tested\n", |
||||
" - Create independent tests that don't rely on each other\n", |
||||
" - Use appropriate fixtures and parametrize when needed\n", |
||||
" - Add clear comments explaining complex test logic\n", |
||||
" - Cover error cases and exceptions\n", |
||||
" - Achieve high code coverage while maintaining meaningful tests\n", |
||||
" \"\"\"\n", |
||||
" return system" |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "22193622-f3a0-4894-a6c4-eb6d88097861", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"def user_prompt_for_comment(code):\n", |
||||
" user = f\"\"\"\n", |
||||
" Please document this Python code with:\n", |
||||
" \n", |
||||
" 1. A docstring containing:\n", |
||||
" - A clear description of purpose and functionality\n", |
||||
" - All parameters with types and descriptions\n", |
||||
" - Return values with types\n", |
||||
" - Exceptions that may be raised\n", |
||||
" - Any important notes or limitations\n", |
||||
" \n", |
||||
" 2. Strategic inline comments for:\n", |
||||
" - Complex algorithms or business logic\n", |
||||
" - Non-obvious implementation choices\n", |
||||
" - Performance considerations\n", |
||||
" - Edge cases\n", |
||||
" \n", |
||||
" Here's the code to document:\n", |
||||
" \\n{code}\n", |
||||
" \"\"\"\n", |
||||
" return user;" |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "81e61752-ec2f-44c1-86a2-ff3234a0358c", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"def user_prompt_for_unit_test(code):\n", |
||||
" user = f\"\"\"\n", |
||||
" Please generate unit tests for the following Python code. Include:\n", |
||||
" \n", |
||||
" 1. Test cases for:\n", |
||||
" - Normal/expected inputs\n", |
||||
" - Edge cases and boundary values\n", |
||||
" - Invalid inputs and error conditions\n", |
||||
" - Different combinations of parameters\n", |
||||
" - All public methods and functions\n", |
||||
" \n", |
||||
" 2. For each test:\n", |
||||
" - Clear test function names describing the scenario\n", |
||||
" - Setup code (fixtures if needed)\n", |
||||
" - Test data preparation\n", |
||||
" - Expected outcomes\n", |
||||
" - Assertions checking results\n", |
||||
" - Comments explaining complex test logic\n", |
||||
" \n", |
||||
" 3. Include any necessary:\n", |
||||
" - Imports\n", |
||||
" - Fixtures\n", |
||||
" - Mock objects\n", |
||||
" - Helper functions\n", |
||||
" - Test data generators\n", |
||||
" \n", |
||||
" Here's the code to test:\n", |
||||
" \\n{code}\n", |
||||
" \"\"\"\n", |
||||
" return user" |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "f31ceed3-0eb2-4962-ab86-2d0302185560", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"pi = \"\"\"\n", |
||||
"import time\n", |
||||
"\n", |
||||
"def calculate(iterations, param1, param2):\n", |
||||
" result = 1.0\n", |
||||
" for i in range(1, iterations+1):\n", |
||||
" j = i * param1 - param2\n", |
||||
" result -= (1/j)\n", |
||||
" j = i * param1 + param2\n", |
||||
" result += (1/j)\n", |
||||
" return result\n", |
||||
"\n", |
||||
"start_time = time.time()\n", |
||||
"result = calculate(100_000_000, 4, 1) * 4\n", |
||||
"end_time = time.time()\n", |
||||
"\n", |
||||
"print(f\"Result: {result:.12f}\")\n", |
||||
"print(f\"Execution Time: {(end_time - start_time):.6f} seconds\")\n", |
||||
"\"\"\"" |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "192c30f5-4be6-49b7-a054-11bfcffa91e0", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"exec(pi)" |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "d4e920dc-4094-42d8-9255-18f2919df2d4", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"def messages_for_comment(python):\n", |
||||
" return [\n", |
||||
" {\"role\": \"system\", \"content\": system_prompt_for_comment()},\n", |
||||
" {\"role\": \"user\", \"content\": user_prompt_for_comment(python)}\n", |
||||
" ]" |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "77500cae-bf84-405c-8b03-2f984108951b", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"def messages_for_unit_test(python):\n", |
||||
" return [\n", |
||||
" {\"role\": \"system\", \"content\": system_prompt_for_unit_test()},\n", |
||||
" {\"role\": \"user\", \"content\": user_prompt_for_unit_test(python)}\n", |
||||
" ]" |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "5ec58bf1-4a44-4c21-a71a-2cac359884e5", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"def stream_comment_gpt(code):\n", |
||||
" stream = openai.chat.completions.create(model=OPENAI_MODEL, messages=messages_for_comment(code), stream=True)\n", |
||||
" reply = \"\"\n", |
||||
" for chunk in stream:\n", |
||||
" fragment = chunk.choices[0].delta.content or \"\"\n", |
||||
" reply += fragment\n", |
||||
" #print(fragment, end='', flush=True)\n", |
||||
" yield reply.replace('```','') \n", |
||||
" " |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "47c615e2-4eb6-4ce1-ad09-7f2e6dbc3934", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"stream_comment_gpt(pi)" |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "0b990875-31fd-40e5-bc8c-f6099d362249", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"def stream_unit_test_gpt(code):\n", |
||||
" stream = openai.chat.completions.create(model=OPENAI_MODEL, messages=messages_for_unit_test(code), stream=True)\n", |
||||
" reply = \"\"\n", |
||||
" for chunk in stream:\n", |
||||
" fragment = chunk.choices[0].delta.content or \"\"\n", |
||||
" reply += fragment\n", |
||||
" #print(fragment, end='', flush=True)\n", |
||||
" yield reply.replace('```','')" |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "3dc90578-4f5e-47f1-b30f-c21b5795e82f", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"stream_unit_test_gpt(pi)" |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "17380c0f-b851-472b-a234-d86f5c219e50", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"def stream_comment_claude(code):\n", |
||||
" result = claude.messages.stream(\n", |
||||
" model=CLAUDE_MODEL,\n", |
||||
" max_tokens=2000,\n", |
||||
" system=system_prompt_for_comment(),\n", |
||||
" messages=[{\"role\": \"user\", \"content\": user_prompt_for_comment(code)}],\n", |
||||
" )\n", |
||||
" reply = \"\"\n", |
||||
" with result as stream:\n", |
||||
" for text in stream.text_stream:\n", |
||||
" reply += text\n", |
||||
" #print(text, end=\"\", flush=True)\n", |
||||
" yield reply.replace('```','')" |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "0a2d016d-76a2-4752-bd4d-6f93ddec46be", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"def stream_unit_test_claude(code):\n", |
||||
" result = claude.messages.stream(\n", |
||||
" model=CLAUDE_MODEL,\n", |
||||
" max_tokens=2000,\n", |
||||
" system=system_prompt_for_unit_test(),\n", |
||||
" messages=[{\"role\": \"user\", \"content\": user_prompt_for_unit_test(code)}],\n", |
||||
" )\n", |
||||
" reply = \"\"\n", |
||||
" with result as stream:\n", |
||||
" for text in stream.text_stream:\n", |
||||
" reply += text\n", |
||||
" #print(text, end=\"\", flush=True)\n", |
||||
" yield reply.replace('```','')" |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "ee43428e-b577-4e95-944d-399f2f3b89ff", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"stream_comment_claude(pi)" |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "0565e33b-9f14-48b7-ae8d-d22dc03b93c9", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"stream_unit_test_claude(pi)" |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "f13b3a5b-366d-4b28-adda-977a313e6b4d", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"def stream_comment_model(model, model_url, code):\n", |
||||
" tokenizer = AutoTokenizer.from_pretrained(model)\n", |
||||
" messages = messages_for_comment(code)\n", |
||||
" text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n", |
||||
" client = InferenceClient(model_url, token=hf_token)\n", |
||||
" stream = client.text_generation(text, stream=True, details=True, max_new_tokens=5000)\n", |
||||
" result = \"\"\n", |
||||
" for r in stream:\n", |
||||
" #print(r.token.text, end = \"\")\n", |
||||
" result += r.token.text\n", |
||||
" yield result \n", |
||||
" " |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "e2efdb92-fc7a-4952-ab46-ae942cb996bf", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"def stream_unit_test_model(model, model_url, code):\n", |
||||
" tokenizer = AutoTokenizer.from_pretrained(model)\n", |
||||
" messages = messages_for_unit_test(code)\n", |
||||
" text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n", |
||||
" client = InferenceClient(model_url, token=hf_token)\n", |
||||
" stream = client.text_generation(text, stream=True, details=True, max_new_tokens=3000)\n", |
||||
" result = \"\"\n", |
||||
" for r in stream:\n", |
||||
" #print(r.token.text, end = \"\")\n", |
||||
" result += r.token.text\n", |
||||
" yield result \n", |
||||
" " |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "0a756193-fcba-43da-a981-203c10d36488", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"stream_comment_model(code_qwen, CODE_QWEN_URL, pi)" |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "12ddcbf4-6286-47a8-847b-5be78e7aa995", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"stream_unit_test_model(code_qwen, CODE_QWEN_URL, pi)" |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "321609ee-b64a-44fc-9090-39f87e1f8e0e", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"def comment_code(python, model):\n", |
||||
" if model==\"GPT\":\n", |
||||
" result = stream_comment_gpt(python)\n", |
||||
" elif model==\"Claude\":\n", |
||||
" result = stream_comment_claude(python)\n", |
||||
" elif model==\"CodeQwen\":\n", |
||||
" result = stream_comment_model(code_qwen, CODE_QWEN_URL, python)\n", |
||||
" else:\n", |
||||
" raise ValueError(\"Unknown model\")\n", |
||||
" for stream_so_far in result:\n", |
||||
" yield stream_so_far " |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "d4c560c9-922d-4893-941f-42893373b1be", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"def get_unit_test(python, model):\n", |
||||
" if model==\"GPT\":\n", |
||||
" result = stream_unit_test_gpt(python)\n", |
||||
" elif model==\"Claude\":\n", |
||||
" result = stream_unit_test_claude(python)\n", |
||||
" elif model==\"CodeQwen\":\n", |
||||
" result = stream_unit_test_model(code_qwen, CODE_QWEN_URL, python)\n", |
||||
" else:\n", |
||||
" raise ValueError(\"Unknown model\")\n", |
||||
" for stream_so_far in result:\n", |
||||
" yield stream_so_far " |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "f85bc777-bebe-436b-88cc-b9ecdb6306c0", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"css = \"\"\"\n", |
||||
".python {background-color: #306998;}\n", |
||||
".cpp {background-color: #050;}\n", |
||||
"\"\"\"" |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "ee27cc91-81e6-42c8-ae3c-c04161229d8c", |
||||
"metadata": { |
||||
"scrolled": true |
||||
}, |
||||
"outputs": [], |
||||
"source": [ |
||||
"with gr.Blocks(css=css) as ui:\n", |
||||
" gr.Markdown(\"## Convert code from Python to C++\")\n", |
||||
" with gr.Row():\n", |
||||
" python = gr.Textbox(label=\"Python code:\", value=pi, lines=10)\n", |
||||
" result = gr.Textbox(label=\"Result code:\", lines=10)\n", |
||||
" with gr.Row():\n", |
||||
" model = gr.Dropdown([\"GPT\", \"Claude\",\"CodeQwen\"], label=\"Select model\", value=\"GPT\")\n", |
||||
" with gr.Row():\n", |
||||
" comment_button = gr.Button(\"Comment code\")\n", |
||||
" with gr.Row():\n", |
||||
" unit_test_button = gr.Button(\"Unit Test code\")\n", |
||||
" \n", |
||||
" comment_button.click(comment_code, inputs=[python, model], outputs=[result])\n", |
||||
" unit_test_button.click(get_unit_test, inputs=[python, model], outputs=[result])\n", |
||||
"ui.launch(inbrowser=True)" |
||||
] |
||||
}, |
||||
{ |
||||
"cell_type": "code", |
||||
"execution_count": null, |
||||
"id": "06e8279c-b488-4807-9bed-9d26be11c057", |
||||
"metadata": {}, |
||||
"outputs": [], |
||||
"source": [] |
||||
} |
||||
], |
||||
"metadata": { |
||||
"kernelspec": { |
||||
"display_name": "Python 3 (ipykernel)", |
||||
"language": "python", |
||||
"name": "python3" |
||||
}, |
||||
"language_info": { |
||||
"codemirror_mode": { |
||||
"name": "ipython", |
||||
"version": 3 |
||||
}, |
||||
"file_extension": ".py", |
||||
"mimetype": "text/x-python", |
||||
"name": "python", |
||||
"nbconvert_exporter": "python", |
||||
"pygments_lexer": "ipython3", |
||||
"version": "3.11.11" |
||||
} |
||||
}, |
||||
"nbformat": 4, |
||||
"nbformat_minor": 5 |
||||
} |
Loading…
Reference in new issue