Update app.py

app.py

@@ -2,6 +2,7 @@ from huggingface_hub import InferenceClient
 import gradio as gr
 from fpdf import FPDF
 from docx import Document
+import os
 
 css = '''
 .gradio-container{max-width: 1000px !important}
@@ -24,20 +25,24 @@ def format_prompt(message, history, system_prompt=None):
     return prompt
 
 def save_to_file(content, filename, format):
-    if format == "pdf":
-        pdf = FPDF()
-        pdf.add_page()
-        pdf.set_auto_page_break(auto=True, margin=15)
-        pdf.set_font("Arial", size=12)
-        pdf.multi_cell(0, 10, content)
-        pdf.output(filename)
-    elif format == "docx":
-        doc = Document()
-        doc.add_paragraph(content)
-        doc.save(filename)
-    elif format == "txt":
-        with open(filename, 'w') as file:
-            file.write(content)
+    try:
+        if format == "pdf":
+            pdf = FPDF()
+            pdf.add_page()
+            pdf.set_auto_page_break(auto=True, margin=15)
+            pdf.set_font("Arial", size=12)
+            pdf.multi_cell(0, 10, content)
+            pdf.output(f"{filename}.pdf")
+        elif format == "docx":
+            doc = Document()
+            doc.add_paragraph(content)
+            doc.save(f"{filename}.docx")
+        elif format == "txt":
+            with open(f"{filename}.txt", 'w') as file:
+                file.write(content)
+        return f"File saved successfully as {filename}.{format}"
+    except Exception as e:
+        return f"Error saving file: {str(e)}"
 
 def generate(
     prompt, history, system_prompt=None, temperature=0.2, max_new_tokens=1024, top_p=0.95, repetition_penalty=1.0,
@@ -66,10 +71,11 @@ def generate(
         output += response.token.text
         yield output
     
-    # Save content to file
     if save_file:
-        save_to_file(output, save_file, save_format)
-    return output
+        # Ensure filename is valid and doesn't contain invalid characters
+        save_file = save_file.strip().replace(' ', '_')
+        save_message = save_to_file(output, save_file, save_format)
+        yield f"{output}\n\n{save_message}"
 
 with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
     gr.Markdown("# GRAB DOC")