Spaces:

SyedTahir
/

STHASI-Generative-AI-Documentation-Generator

Sleeping

STHASI-Generative-AI-Documentation-Generator / app.py

Create app.py

f05e9f6 verified about 1 year ago

10.4 kB

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"""
    <s>[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"""
    <s>[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)

	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"""
	<s>[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"""
	<s>[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)