Update app.py

app.py

@@ -1,21 +1,17 @@
 import gradio as gr
 from openai import OpenAI
 import os
-from io import BytesIO
-from reportlab.lib.pagesizes import letter
-from reportlab.pdfgen import canvas
+from fpdf import FPDF
 from docx import Document
 
-# Custom CSS
 css = '''
-.gradio-container{max-width: 1500px !important}
+.gradio-container{max-width: 1000px !important}
 h1{text-align:center}
 footer {
     visibility: hidden
 }
 '''
 
-# Set up OpenAI client
 ACCESS_TOKEN = os.getenv("HF_TOKEN")
 
 client = OpenAI(
@@ -23,7 +19,6 @@ client = OpenAI(
     api_key=ACCESS_TOKEN,
 )
 
-# Function to handle chat responses
 def respond(
     message,
     history: list[tuple[str, str]],
@@ -44,7 +39,7 @@ def respond(
 
     response = ""
     
-    for message in client.chat.completions.create(
+    for message in  client.chat.completions.create(
         model="meta-llama/Meta-Llama-3.1-8B-Instruct",
         max_tokens=max_tokens,
         stream=True,
@@ -53,99 +48,53 @@ def respond(
         messages=messages,
     ):
         token = message.choices[0].delta.content
+        
         response += token
         yield response
 
-# Function to save chat history to a text file
-def save_as_txt(history):
-    file_path = "chat_history.txt"
-    with open(file_path, "w") as f:
+def save_to_file(history, 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 user_message, assistant_message in history:
-            f.write(f"User: {user_message}\n")
-            f.write(f"Assistant: {assistant_message}\n")
-    return file_path
-
-# Function to save chat history to a DOCX file
-def save_as_docx(history):
-    file_path = "chat_history.docx"
-    doc = Document()
-    doc.add_heading('Chat History', 0)
-    
-    for user_message, assistant_message in history:
-        doc.add_paragraph(f"User: {user_message}")
-        doc.add_paragraph(f"Assistant: {assistant_message}")
-    
-    doc.save(file_path)
-    return file_path
-
-# Function to save chat history to a PDF file
-def save_as_pdf(history):
-    file_path = "chat_history.pdf"
-    buffer = BytesIO()
-    c = canvas.Canvas(buffer, pagesize=letter)
-    width, height = letter
-    y = height - 40
-    
-    c.drawString(30, y, "Chat History")
-    y -= 30
-    
-    for user_message, assistant_message in history:
-        c.drawString(30, y, f"User: {user_message}")
-        y -= 20
-        c.drawString(30, y, f"Assistant: {assistant_message}")
-        y -= 30
+            pdf.multi_cell(0, 10, f"User: {user_message}")
+            pdf.multi_cell(0, 10, f"Assistant: {assistant_message}")
+        file_name = "chat_history.pdf"
+        pdf.output(file_name)
         
-        if y < 40:
-            c.showPage()
-            y = height - 40
-            
-    c.save()
-    buffer.seek(0)
-    
-    with open(file_path, "wb") as f:
-        f.write(buffer.read())
-    
-    return file_path
-
-# Function to handle file saving based on format
-def handle_file_save(history, file_format):
-    if file_format == "txt":
-        return save_as_txt(history)
-    elif file_format == "docx":
-        return save_as_docx(history)
-    elif file_format == "pdf":
-        return save_as_pdf(history)
-    return None
+    elif file_format == "DOCX":
+        doc = Document()
+        for user_message, assistant_message in history:
+            doc.add_paragraph(f"User: {user_message}")
+            doc.add_paragraph(f"Assistant: {assistant_message}")
+        file_name = "chat_history.docx"
+        doc.save(file_name)
+        
+    elif file_format == "TXT":
+        file_name = "chat_history.txt"
+        with open(file_name, "w") as file:
+            for user_message, assistant_message in history:
+                file.write(f"User: {user_message}\n")
+                file.write(f"Assistant: {assistant_message}\n")
+                
+    return file_name
 
-# Handler function for Gradio app
-def save_handler(message, history, system_message, max_tokens, temperature, top_p, file_format):
-    new_history = history + [(message, next(respond(message, history, system_message, max_tokens, temperature, top_p)))]
-    saved_file = handle_file_save(new_history, file_format)
-    return saved_file, new_history
+def save_conversation(history, file_format):
+    file_name = save_to_file(history, file_format)
+    return file_name
 
-# Gradio interface
-demo = gr.Interface(
-    fn=save_handler,
-    inputs=[
-        gr.Textbox(value="", label="Type a message..", lines=5),
-        gr.State([]),  # Initialize state as an empty list
-        gr.Textbox(value="", label="System message", visible=False),
+demo = gr.ChatInterface(
+    respond,
+    additional_inputs=[
+        gr.Textbox(value="", label="System message"),
         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-P",
-        ),
-        gr.Dropdown(
-            choices=["txt", "docx", "pdf"],
-            label="Save as",
-            value="pdf",
-        ),
+        gr.Slider(minimum=0.1, maximum 1.0, value=0.95, step=0.05, label="Top-P"),
+        gr.Radio(["PDF", "DOCX", "TXT"], label="Save As"),
     ],
-    outputs=[gr.File(label="Download Chat History"), gr.State()],
+    button_fn=save_conversation,
     css=css,
     theme="allenai/gradio-theme",
 )