|
import os |
|
from pathlib import Path |
|
|
|
import gradio as gr |
|
import torch |
|
from huggingface_hub import snapshot_download |
|
from transformers import pipeline |
|
|
|
model_name = "databricks/dolly-v2-12b" |
|
|
|
local_dir = f"./models/{model_name}" |
|
if not Path(local_dir).exists() or len(os.listdir(local_dir)) == 0: |
|
snapshot_download(model_name, local_dir=local_dir) |
|
|
|
generate_text = pipeline(model=local_dir, torch_dtype=torch.bfloat16, trust_remote_code=True, |
|
device_map="auto") |
|
|
|
theme = gr.themes.Monochrome( |
|
primary_hue="indigo", |
|
secondary_hue="blue", |
|
neutral_hue="slate", |
|
radius_size=gr.themes.sizes.radius_sm, |
|
font=[gr.themes.GoogleFont("Open Sans"), "ui-sans-serif", "system-ui", "sans-serif"], |
|
) |
|
|
|
with gr.Blocks(theme=theme) as demo: |
|
chatbot = gr.Chatbot() |
|
msg = gr.Textbox() |
|
clear = gr.Button("Clear") |
|
|
|
|
|
def user(user_message, history): |
|
return "", history + [[user_message, None]] |
|
|
|
|
|
def bot(history): |
|
history[-1][1] = generate_text(history[-1][0]) |
|
return history |
|
|
|
|
|
msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then( |
|
bot, chatbot, chatbot |
|
) |
|
clear.click(lambda: None, None, chatbot, queue=False) |
|
|
|
if __name__ == "__main__": |
|
demo.launch(server_name="0.0.0.0", server_port=3000) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
