Create app.py

app.py

@@ -0,0 +1,277 @@
+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)