SyedTahir/STHASI-Generative-AI-Documentation-Generator
0
1import gradio as gr2import os3import git4import ast5import torch6from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig7from threading import Thread8 9# --- 1. Configuration and Model Loading ---10 11# Get the Hugging Face token from environment variables12HF_TOKEN = os.getenv("HF_TOKEN")13 14# Define the model ID for CodeLlama-7b-Instruct15MODEL_ID = "codellama/CodeLlama-7b-Instruct-hf"16 17# --- Model Loading (with optimizations for Spaces) ---18# Use 4-bit quantization to reduce memory usage19bnb_config = BitsAndBytesConfig(20 load_in_4bit=True,21 bnb_4bit_quant_type="nf4",22 bnb_4bit_compute_dtype=torch.float16,23)24 25# Load the tokenizer and model26tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_auth_token=HF_TOKEN)27model = AutoModelForCausalLM.from_pretrained(28 MODEL_ID,29 quantization_config=bnb_config,30 device_map="auto", # Automatically maps model layers to available hardware31 use_auth_token=HF_TOKEN,32)33 34# --- 2. Backend Logic: Code Parsing and Generation ---35 36def parse_python_code(file_path):37 """38 Parses a Python file to extract functions and classes using AST.39 This is more robust than simple text parsing.40 """41 with open(file_path, 'r', encoding='utf-8') as file:42 code = file.read()43 44 tree = ast.parse(code)45 elements = []46 for node in ast.walk(tree):47 if isinstance(node, ast.FunctionDef):48 elements.append({49 "type": "Function",50 "name": node.name,51 "code": ast.get_source_segment(code, node)52 })53 elif isinstance(node, ast.ClassDef):54 elements.append({55 "type": "Class",56 "name": node.name,57 "code": ast.get_source_segment(code, node)58 })59 return elements60 61def clone_repository(repo_url):62 """63 Clones a public GitHub repository to a local directory.64 The directory is named after the repository.65 """66 try:67 repo_name = repo_url.split('/')[-1].replace('.git', '')68 clone_path = os.path.join("/tmp", repo_name) # Use /tmp for temporary storage in Spaces69 70 if os.path.exists(clone_path):71 # If exists, pull latest changes72 repo = git.Repo(clone_path)73 repo.remotes.origin.pull()74 else:75 # Otherwise, clone it76 git.Repo.clone_from(repo_url, clone_path)77 78 return clone_path, None79 except Exception as e:80 return None, str(e)81 82def generate_documentation(code_snippet):83 """84 Generates documentation for a single code snippet using CodeLlama.85 """86 prompt = f"""87 <s>[INST] You are an expert programmer tasked with writing clear and concise documentation for Python code. 88 Generate a high-quality docstring for the following code snippet.89 90 The documentation should explain:91 1. The purpose of the function/class.92 2. The arguments it takes (if any), including their types and what they do.93 3. What it returns (if anything).94 4. A simple usage example in a code block.95 96 Do not repeat the code in your response, only provide the docstring content.97 98 Code:99 ```python100 {code_snippet}101 ```102 [/INST]103 """104 105 inputs = tokenizer(prompt, return_tensors="pt").to("cuda") # Ensure tensors are on the GPU106 107 # Use a streaming generator for better UX108 streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)109 generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=512)110 111 thread = Thread(target=model.generate, kwargs=generation_kwargs)112 thread.start()113 114 # Yield generated text as it comes115 for new_text in streamer:116 yield new_text117 118def process_repository(repo_url):119 """120 Main function to clone, parse, and generate documentation for a repository.121 Yields updates to the Gradio interface.122 """123 yield "Cloning repository...", None, gr.update(visible=True)124 125 clone_path, error = clone_repository(repo_url)126 if error:127 yield f"Error cloning repository: {error}", None, gr.update(visible=False)128 return129 130 yield f"Repository cloned to {clone_path}. Parsing Python files...", None, gr.update(visible=True)131 132 all_docs = ""133 # Walk through all directories and files in the cloned repo134 for root, _, files in os.walk(clone_path):135 for file in files:136 if file.endswith('.py'):137 file_path = os.path.join(root, file)138 relative_path = os.path.relpath(file_path, clone_path)139 140 yield f"Parsing `{relative_path}`...", all_docs, gr.update(visible=True)141 142 code_elements = parse_python_code(file_path)143 if not code_elements:144 continue145 146 all_docs += f"\n\n---\n\n### 📄 **File: `{relative_path}`**\n\n"147 148 for element in code_elements:149 all_docs += f"#### 🧬 {element['type']}: `{element['name']}`\n\n"150 all_docs += "```python\n"151 all_docs += f"{element['code']}\n"152 all_docs += "```\n\n"153 all_docs += "**🤖 AI-Generated Documentation:**\n"154 155 # Stream the documentation for this element156 full_doc_string = ""157 for new_text in generate_documentation(element['code']):158 full_doc_string += new_text159 yield f"Generating docs for `{element['name']}` in `{relative_path}`...", all_docs + full_doc_string, gr.update(visible=True)160 all_docs += full_doc_string161 162 yield "Documentation generation complete!", all_docs, gr.update(visible=True)163 164 165# --- 3. Advanced Feature: "Ask Your Codebase" Chatbot ---166 167def get_codebase_context(clone_path):168 """Gathers all python code from the repo into a single string for context."""169 context = ""170 for root, _, files in os.walk(clone_path):171 for file in files:172 if file.endswith('.py'):173 file_path = os.path.join(root, file)174 with open(file_path, 'r', encoding='utf-8') as f:175 context += f"--- File: {os.path.relpath(file_path, clone_path)} ---\n{f.read()}\n\n"176 return context177 178def answer_question(repo_url, question, history):179 """Answers a user's question based on the codebase context."""180 if not repo_url:181 return "Please provide a repository URL first.", history182 183 clone_path, _ = clone_repository(repo_url)184 if not clone_path:185 return "Could not clone the repository to answer the question.", history186 187 code_context = get_codebase_context(clone_path)188 189 # Limit context size to avoid exceeding model limits190 max_context_length = 10000 191 if len(code_context) > max_context_length:192 code_context = code_context[:max_context_length] + "\n... (context truncated)"193 194 prompt = f"""195 <s>[INST] You are a helpful AI assistant that answers questions about a software project.196 Use the following codebase context to answer the user's question.197 If the answer is not in the context, say that you cannot find the answer in the provided code.198 199 --- Codebase Context ---200 {code_context}201 202 --- Question ---203 {question}204 [/INST]205 """206 207 inputs = tokenizer(prompt, return_tensors="pt").to("cuda")208 outputs = model.generate(inputs, max_new_tokens=1024, temperature=0.7)209 answer = tokenizer.decode(outputs[0], skip_special_tokens=True)210 211 # Clean up the output to only return the answer part212 answer_start = answer.find("[/INST]") + len("[/INST]")213 final_answer = answer[answer_start:].strip()214 215 history.append((question, final_answer))216 return "", history217 218# --- 4. Gradio UI ---219 220# Custom CSS for a modern, dark theme221custom_css = """222body { background-color: #1a1a1a; color: #f0f0f0; }223.gradio-container { max-width: 90% !important; margin: auto; }224.gr-button { background-color: #3498db; color: white; border: none; }225.gr-button:hover { background-color: #2980b9; }226#status_textbox { text-align: center; color: #a0a0a0; font-style: italic; }227#doc_output .prose { color: #f0f0f0; }228footer { display: none !important; }229"""230 231with gr.Blocks(theme=gr.themes.Base(), css=custom_css) as demo:232 gr.Markdown("# 🚀 Generative AI Documentation Generator")233 gr.Markdown("Enter a public GitHub repository URL to automatically generate documentation for its Python code.")234 235 with gr.Tabs():236 with gr.TabItem("📄 Documentation Generator"):237 with gr.Row():238 repo_url_input = gr.Textbox(239 label="GitHub Repository URL", 240 placeholder="e.g., https://github.com/gradio-app/gradio"241 )242 generate_button = gr.Button("✨ Generate Docs", variant="primary")243 244 status_textbox = gr.Textbox(245 label="Status", 246 interactive=False, 247 elem_id="status_textbox"248 )249 250 doc_output = gr.Markdown(elem_id="doc_output")251 252 # Connect the button to the processing function253 generate_button.click(254 fn=process_repository,255 inputs=[repo_url_input],256 outputs=[status_textbox, doc_output]257 )258 259 with gr.TabItem("💬 Ask Your Codebase (Advanced)"):260 repo_url_chatbot = gr.Textbox(261 label="GitHub Repository URL (must be the same as used for generation)", 262 placeholder="Enter the repo URL here to enable the chatbot"263 )264 chatbot_ui = gr.Chatbot(label="Chat About Your Code")265 question_input = gr.Textbox(label="Your Question", placeholder="e.g., 'What does the 'predict' function do?'")266 submit_button = gr.Button("Submit Question")267 268 # Connect the submit button to the chatbot function269 submit_button.click(270 fn=answer_question,271 inputs=[repo_url_chatbot, question_input, chatbot_ui],272 outputs=[question_input, chatbot_ui]273 )274 275if __name__ == "__main__":276 # Launch the Gradio app with sharing enabled for easy access277 demo.launch(share=True)