GRAB-DOC / app.py
prithivMLmods's picture
Update app.py
68c5b4c verified
raw
history blame
No virus
3.55 kB
import gradio as gr
from openai import OpenAI
import os
from fpdf import FPDF
from docx import Document
css = '''
.gradio-container{max-width: 1000px !important}
h1{text-align:center}
footer {
visibility: hidden
}
'''
ACCESS_TOKEN = os.getenv("HF_TOKEN")
client = OpenAI(
base_url="https://api-inference.huggingface.co/v1/",
api_key=ACCESS_TOKEN,
)
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
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:
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)
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
# Gradio Interface Setup
with gr.Blocks(css=css) as demo:
system_message = gr.Textbox(value="", label="System message")
max_tokens = gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens")
temperature = gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature")
top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-P")
save_as = gr.Radio(["PDF", "DOCX", "TXT"], label="Save As")
chat = gr.Chatbot()
msg = gr.Textbox(label="Your message")
def respond_wrapper(message, history):
response_generator = respond(
message,
history,
system_message.value,
max_tokens.value,
temperature.value,
top_p.value
)
response = next(response_generator)
return history + [(message, response)]
msg.submit(respond_wrapper, [msg, chat], [chat])
save_button = gr.Button("Save Conversation")
output_file = gr.File(label="Download File")
def handle_save(history, file_format):
return save_to_file(history, file_format)
save_button.click(handle_save, inputs=[chat, save_as], outputs=output_file)
if __name__ == "__main__":
demo.launch()