Spaces:
Sleeping
Sleeping
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() |