Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import spaces
+import io
 
 
 # Initialize the Qwen model and tokenizer
@@ -12,7 +13,7 @@ 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}"
+    prompt = f"Generate README 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."},
@@ -34,7 +35,11 @@ def generate_documentation(code_input):
 # Gradio interface
 def process_code(code_input):
     documentation = generate_documentation(code_input)
-    return documentation
+    readme_file = io.BytesIO()
+    readme_file.write(documentation.encode('utf-8'))
+    readme_file.seek(0)
+    
+    return documentation, readme_file
 
 # 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:
@@ -51,10 +56,11 @@ with gr.Blocks(css=".container { font-family: 'Roboto', sans-serif; } .btn-prima
         generate_button = gr.Button("Generate Documentation", elem_classes="btn btn-primary")
     
     with gr.Row():
-        output_text = gr.Markdown()
-    
+        output_text = gr.Markdown(label="Generated Documentation")
+        download_button = gr.File(label="Download README.md")
+
     # Bind function to button click
-    generate_button.click(process_code, inputs=code_input, outputs=output_text)
+    generate_button.click(process_code, inputs=code_input, outputs=[output_text, download_button])
 
 # Launch the Gradio app
 app.launch()