import gradio as gr import os import git import ast import torch from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig from threading import Thread # --- 1. Configuration and Model Loading --- # Get the Hugging Face token from environment variables HF_TOKEN = os.getenv("HF_TOKEN") # Define the model ID for CodeLlama-7b-Instruct MODEL_ID = "codellama/CodeLlama-7b-Instruct-hf" # --- Model Loading (with optimizations for Spaces) --- # Use 4-bit quantization to reduce memory usage bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.float16, ) # Load the tokenizer and model tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_auth_token=HF_TOKEN) model = AutoModelForCausalLM.from_pretrained( MODEL_ID, quantization_config=bnb_config, device_map="auto", # Automatically maps model layers to available hardware use_auth_token=HF_TOKEN, ) # --- 2. Backend Logic: Code Parsing and Generation --- def parse_python_code(file_path): """ Parses a Python file to extract functions and classes using AST. This is more robust than simple text parsing. """ with open(file_path, 'r', encoding='utf-8') as file: code = file.read() tree = ast.parse(code) elements = [] for node in ast.walk(tree): if isinstance(node, ast.FunctionDef): elements.append({ "type": "Function", "name": node.name, "code": ast.get_source_segment(code, node) }) elif isinstance(node, ast.ClassDef): elements.append({ "type": "Class", "name": node.name, "code": ast.get_source_segment(code, node) }) return elements def clone_repository(repo_url): """ Clones a public GitHub repository to a local directory. The directory is named after the repository. """ try: repo_name = repo_url.split('/')[-1].replace('.git', '') clone_path = os.path.join("/tmp", repo_name) # Use /tmp for temporary storage in Spaces if os.path.exists(clone_path): # If exists, pull latest changes repo = git.Repo(clone_path) repo.remotes.origin.pull() else: # Otherwise, clone it git.Repo.clone_from(repo_url, clone_path) return clone_path, None except Exception as e: return None, str(e) def generate_documentation(code_snippet): """ Generates documentation for a single code snippet using CodeLlama. """ prompt = f""" [INST] You are an expert programmer tasked with writing clear and concise documentation for Python code. Generate a high-quality docstring for the following code snippet. The documentation should explain: 1. The purpose of the function/class. 2. The arguments it takes (if any), including their types and what they do. 3. What it returns (if anything). 4. A simple usage example in a code block. Do not repeat the code in your response, only provide the docstring content. Code: ```python {code_snippet} ``` [/INST] """ inputs = tokenizer(prompt, return_tensors="pt").to("cuda") # Ensure tensors are on the GPU # Use a streaming generator for better UX streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=512) thread = Thread(target=model.generate, kwargs=generation_kwargs) thread.start() # Yield generated text as it comes for new_text in streamer: yield new_text def process_repository(repo_url): """ Main function to clone, parse, and generate documentation for a repository. Yields updates to the Gradio interface. """ yield "Cloning repository...", None, gr.update(visible=True) clone_path, error = clone_repository(repo_url) if error: yield f"Error cloning repository: {error}", None, gr.update(visible=False) return yield f"Repository cloned to {clone_path}. Parsing Python files...", None, gr.update(visible=True) all_docs = "" # Walk through all directories and files in the cloned repo for root, _, files in os.walk(clone_path): for file in files: if file.endswith('.py'): file_path = os.path.join(root, file) relative_path = os.path.relpath(file_path, clone_path) yield f"Parsing `{relative_path}`...", all_docs, gr.update(visible=True) code_elements = parse_python_code(file_path) if not code_elements: continue all_docs += f"\n\n---\n\n### 📄 **File: `{relative_path}`**\n\n" for element in code_elements: all_docs += f"#### 🧬 {element['type']}: `{element['name']}`\n\n" all_docs += "```python\n" all_docs += f"{element['code']}\n" all_docs += "```\n\n" all_docs += "**🤖 AI-Generated Documentation:**\n" # Stream the documentation for this element full_doc_string = "" for new_text in generate_documentation(element['code']): full_doc_string += new_text yield f"Generating docs for `{element['name']}` in `{relative_path}`...", all_docs + full_doc_string, gr.update(visible=True) all_docs += full_doc_string yield "Documentation generation complete!", all_docs, gr.update(visible=True) # --- 3. Advanced Feature: "Ask Your Codebase" Chatbot --- def get_codebase_context(clone_path): """Gathers all python code from the repo into a single string for context.""" context = "" for root, _, files in os.walk(clone_path): for file in files: if file.endswith('.py'): file_path = os.path.join(root, file) with open(file_path, 'r', encoding='utf-8') as f: context += f"--- File: {os.path.relpath(file_path, clone_path)} ---\n{f.read()}\n\n" return context def answer_question(repo_url, question, history): """Answers a user's question based on the codebase context.""" if not repo_url: return "Please provide a repository URL first.", history clone_path, _ = clone_repository(repo_url) if not clone_path: return "Could not clone the repository to answer the question.", history code_context = get_codebase_context(clone_path) # Limit context size to avoid exceeding model limits max_context_length = 10000 if len(code_context) > max_context_length: code_context = code_context[:max_context_length] + "\n... (context truncated)" prompt = f""" [INST] You are a helpful AI assistant that answers questions about a software project. Use the following codebase context to answer the user's question. If the answer is not in the context, say that you cannot find the answer in the provided code. --- Codebase Context --- {code_context} --- Question --- {question} [/INST] """ inputs = tokenizer(prompt, return_tensors="pt").to("cuda") outputs = model.generate(inputs, max_new_tokens=1024, temperature=0.7) answer = tokenizer.decode(outputs[0], skip_special_tokens=True) # Clean up the output to only return the answer part answer_start = answer.find("[/INST]") + len("[/INST]") final_answer = answer[answer_start:].strip() history.append((question, final_answer)) return "", history # --- 4. Gradio UI --- # Custom CSS for a modern, dark theme custom_css = """ body { background-color: #1a1a1a; color: #f0f0f0; } .gradio-container { max-width: 90% !important; margin: auto; } .gr-button { background-color: #3498db; color: white; border: none; } .gr-button:hover { background-color: #2980b9; } #status_textbox { text-align: center; color: #a0a0a0; font-style: italic; } #doc_output .prose { color: #f0f0f0; } footer { display: none !important; } """ with gr.Blocks(theme=gr.themes.Base(), css=custom_css) as demo: gr.Markdown("# 🚀 Generative AI Documentation Generator") gr.Markdown("Enter a public GitHub repository URL to automatically generate documentation for its Python code.") with gr.Tabs(): with gr.TabItem("📄 Documentation Generator"): with gr.Row(): repo_url_input = gr.Textbox( label="GitHub Repository URL", placeholder="e.g., https://github.com/gradio-app/gradio" ) generate_button = gr.Button("✨ Generate Docs", variant="primary") status_textbox = gr.Textbox( label="Status", interactive=False, elem_id="status_textbox" ) doc_output = gr.Markdown(elem_id="doc_output") # Connect the button to the processing function generate_button.click( fn=process_repository, inputs=[repo_url_input], outputs=[status_textbox, doc_output] ) with gr.TabItem("💬 Ask Your Codebase (Advanced)"): repo_url_chatbot = gr.Textbox( label="GitHub Repository URL (must be the same as used for generation)", placeholder="Enter the repo URL here to enable the chatbot" ) chatbot_ui = gr.Chatbot(label="Chat About Your Code") question_input = gr.Textbox(label="Your Question", placeholder="e.g., 'What does the 'predict' function do?'") submit_button = gr.Button("Submit Question") # Connect the submit button to the chatbot function submit_button.click( fn=answer_question, inputs=[repo_url_chatbot, question_input, chatbot_ui], outputs=[question_input, chatbot_ui] ) if __name__ == "__main__": # Launch the Gradio app with sharing enabled for easy access demo.launch(share=True)