import gradio as gr from transformers import AutoModelForCausalLM, AutoTokenizer from fpdf import FPDF import torch import spaces # Initialize the Qwen model and tokenizer model_name = "Qwen/Qwen2.5-Coder-7B-Instruct" model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto", device_map="auto") tokenizer = AutoTokenizer.from_pretrained(model_name) # Function to generate README and documentation @spaces.GPU def generate_documentation(code_input): prompt = f"Generate README and documentation for the following code:\n\n{code_input}" messages = [ {"role": "system", "content": "You are CodeDocify, a highly efficient and intelligent assistant designed to analyze code and generate comprehensive, clear, and concise documentation. Your purpose is to help developers by producing well-structured README files and detailed explanations of their code. You aim to simplify complex code into easily understandable documentation, ensuring that your responses are accurate, professional, and easy to follow."}, {"role": "user", "content": prompt} ] # Prepare inputs for the model text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) model_inputs = tokenizer([text], return_tensors="pt").to(model.device) # Generate the documentation generated_ids = model.generate(**model_inputs, max_new_tokens=4000) generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)] documentation = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0] print(documentation) return documentation # Function to generate and download PDF def create_pdf(documentation): pdf = FPDF() pdf.set_auto_page_break(auto=True, margin=15) pdf.add_page() pdf.set_font("Arial", size=12) pdf.multi_cell(200, 10, documentation) # Create a bytes buffer and output the PDF to it buffer = io.BytesIO() pdf.output(buffer) buffer.seek(0) return buffer # Gradio interface def process_code(code_input): documentation = generate_documentation(code_input) pdf_path = create_pdf(documentation) return documentation, pdf_path # Set up the Gradio app with Bootstrap, icons, and smiley with gr.Blocks(css=".container { font-family: 'Roboto', sans-serif; } .btn-primary { background-color: #007bff; } .icon { margin-right: 10px; }") as app: gr.Markdown(""" #Code Documentation Generator Paste your code below, and the app will generate the README and detailed documentation for you. The output will also be available for download as a PDF. """) with gr.Row(): code_input = gr.Textbox(lines=10, label="Paste your code here", placeholder="Enter your code...", show_label=False, elem_classes="form-control") with gr.Row(): generate_button = gr.Button("Generate Documentation", elem_classes="btn btn-primary") with gr.Row(): output_text = gr.Textbox(label="Generated Documentation", lines=20, interactive=False) download_pdf = gr.File(label="Download PDF", file_types=[".pdf"]) # Bind function to button click generate_button.click(process_code, inputs=code_input, outputs=[output_text, download_pdf]) # Launch the Gradio app app.launch()