Update app.py

app.py

@@ -1,5 +1,7 @@
 from huggingface_hub import InferenceClient
 import gradio as gr
+from fpdf import FPDF
+import docx
 
 css = '''
 .gradio-container{max-width: 1000px !important}
@@ -10,6 +12,7 @@ footer {
 '''
 
 client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3")
+
 def format_prompt(message, history, system_prompt=None):
     prompt = "<s>"
     for user_prompt, bot_response in history:
@@ -19,7 +22,8 @@ def format_prompt(message, history, system_prompt=None):
         prompt += f"[SYS] {system_prompt} [/SYS]"
     prompt += f"[INST] {message} [/INST]"
     return prompt
-#Generate
+
+# Generate text
 def generate(
     prompt, history, system_prompt=None, temperature=0.2, max_new_tokens=1024, top_p=0.95, repetition_penalty=1.0,
 ):
@@ -47,9 +51,50 @@ def generate(
         yield output
     return output
 
+# Save the generated content to a file
+def save_file(content, filename, file_format):
+    if file_format == "pdf":
+        pdf = FPDF()
+        pdf.add_page()
+        pdf.set_auto_page_break(auto=True, margin=15)
+        pdf.set_font("Arial", size=12)
+        for line in content.split("\n"):
+            pdf.multi_cell(0, 10, line)
+        pdf.output(f"{filename}.pdf")
+        return f"{filename}.pdf"
+    elif file_format == "docx":
+        doc = docx.Document()
+        doc.add_paragraph(content)
+        doc.save(f"{filename}.docx")
+        return f"{filename}.docx"
+    elif file_format == "txt":
+        with open(f"{filename}.txt", "w") as f:
+            f.write(content)
+        return f"{filename}.txt"
+    else:
+        raise ValueError("Unsupported file format")
+
+# Combine generate and save file functions
+def generate_and_save(prompt, history, filename="output", file_format="pdf", system_prompt=None, temperature=0.2, max_new_tokens=1024, top_p=0.95, repetition_penalty=1.0):
+    generated_text = ""
+    for output in generate(prompt, history, system_prompt, temperature, max_new_tokens, top_p, repetition_penalty):
+        generated_text = output
+    saved_file = save_file(generated_text, filename, file_format)
+    return generated_text, saved_file
 
-demo = gr.ChatInterface(
-    fn=generate,
+# Create Gradio interface
+demo = gr.Interface(
+    fn=generate_and_save,
+    inputs=[
+        gr.Textbox(placeholder="Type your message here...", label="Prompt"),
+        gr.State(value=[]),  # history
+        gr.Textbox(placeholder="Filename (default: output)", label="Filename", value="output"),
+        gr.Radio(["pdf", "docx", "txt"], label="File Format", value="pdf"),
+    ],
+    outputs=[
+        gr.Textbox(label="Generated Text"),
+        gr.File(label="Download File")
+    ],
     css=css,
     title="",
     theme="bethecloud/storj_theme"