GRAB-DOC / app.py
prithivMLmods's picture
Update app.py
edb2d41 verified
raw
history blame
No virus
4.68 kB
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 docx import Document
# Custom CSS
css = '''
.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(
base_url="https://api-inference.huggingface.co/v1/",
api_key=ACCESS_TOKEN,
)
# Function to handle chat responses
def respond(
message,
history: list[tuple[str, str]],
system_message,
max_tokens,
temperature,
top_p,
):
messages = [{"role": "system", "content": system_message}]
for val in history:
if val[0]:
messages.append({"role": "user", "content": val[0]})
if val[1]:
messages.append({"role": "assistant", "content": val[1]})
messages.append({"role": "user", "content": message})
response = ""
for message in client.chat.completions.create(
model="meta-llama/Meta-Llama-3.1-8B-Instruct",
max_tokens=max_tokens,
stream=True,
temperature=temperature,
top_p=top_p,
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):
with open("chat_history.txt", "w") as f:
for user_message, assistant_message in history:
f.write(f"User: {user_message}\n")
f.write(f"Assistant: {assistant_message}\n")
return "chat_history.txt"
# Function to save chat history to a DOCX file
def save_as_docx(history):
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("chat_history.docx")
return "chat_history.docx"
# Function to save chat history to a PDF file
def save_as_pdf(history):
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
if y < 40:
c.showPage()
y = height - 40
c.save()
buffer.seek(0)
with open("chat_history.pdf", "wb") as f:
f.write(buffer.read())
return "chat_history.pdf"
# Gradio interface
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)
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",
),
],
outputs=[
gr.File(label="Download Chat History"),
],
css=css,
theme="allenai/gradio-theme",
)
def save_handler(message, history, system_message, max_tokens, temperature, top_p, file_format):
response = respond(message, history, system_message, max_tokens, temperature, top_p)
saved_file = handle_file_save(history, file_format)
return saved_file
demo = gr.Interface(
fn=save_handler,
inputs=[
gr.Textbox(value="", label="Message"),
gr.State(),
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",
),
],
outputs=gr.File(label="Download Chat History"),
css=css,
theme="allenai/gradio-theme",
)
if __name__ == "__main__":
demo.launch()