Spaces:
Sleeping
Sleeping
| # app.py ────────────────────────────────────────────────────────────────────── | |
| import torch | |
| import gradio as gr | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| TITLE = "Talk To Me Morty" | |
| DESCRIPTION = """ | |
| <p style='text-align:center'> | |
| The bot was trained on a Rick & Morty dialogues dataset with DialoGPT. | |
| </p> | |
| <center> | |
| <img src="https://huggingface.co/spaces/kingabzpro/Rick_and_Morty_Bot/resolve/main/img/rick.png" | |
| alt="Rick" | |
| width="150"> | |
| </center> | |
| """ | |
| ARTICLE = """ | |
| <p style='text-align:center'> | |
| <a href="https://medium.com/geekculture/discord-bot-using-dailogpt-and-huggingface-api-c71983422701" | |
| target="_blank">Complete Tutorial</a> · | |
| <a href="https://dagshub.com/kingabzpro/DailoGPT-RickBot" | |
| target="_blank">Project on DAGsHub</a> | |
| </p> | |
| """ | |
| # ─── Load model once at start ──────────────────────────────────────────────── | |
| tokenizer = AutoTokenizer.from_pretrained("ericzhou/DialoGPT-Medium-Rick_v2") | |
| model = AutoModelForCausalLM.from_pretrained("ericzhou/DialoGPT-Medium-Rick_v2") | |
| # ─── Chat handler ──────────────────────────────────────────────────────────── | |
| def chat(user_msg: str, history_ids: list[int] | None): | |
| if not user_msg: | |
| return [], history_ids or [] | |
| new_ids = tokenizer.encode(user_msg + tokenizer.eos_token, | |
| return_tensors="pt") | |
| bot_input = ( | |
| torch.cat([torch.LongTensor(history_ids), new_ids], dim=-1) | |
| if history_ids else new_ids | |
| ) | |
| history_ids = model.generate( | |
| bot_input, | |
| max_length=min(4096, bot_input.shape[-1] + 200), | |
| pad_token_id=tokenizer.eos_token_id, | |
| do_sample=True, | |
| temperature=0.7, | |
| top_p=0.92, | |
| top_k=50, | |
| ).tolist() | |
| turns = tokenizer.decode(history_ids[0], skip_special_tokens=False) \ | |
| .split("<|endoftext|>") | |
| # pack into (user, bot) pairs for Chatbot component | |
| pairs = [(turns[i], turns[i + 1]) for i in range(0, len(turns) - 1, 2)] | |
| return pairs, history_ids | |
| # ─── Gradio UI ─────────────────────────────────────────────────────────────── | |
| with gr.Blocks(theme=gr.themes.Base()) as demo: | |
| gr.Markdown(f"<h1 style='text-align:center'>{TITLE}</h1>") | |
| gr.Markdown(DESCRIPTION) | |
| chatbot = gr.Chatbot(height=450) | |
| state = gr.State([]) | |
| with gr.Row(equal_height=True): | |
| prompt = gr.Textbox(placeholder="Ask Rick anything…", scale=9, show_label=False) | |
| send = gr.Button("Send",scale=1, variant="primary") | |
| # send on click or ↵ | |
| send.click(chat, inputs=[prompt, state], outputs=[chatbot, state]) | |
| prompt.submit(chat, inputs=[prompt, state], outputs=[chatbot, state]) | |
| gr.Examples([["How are you, Rick?"], ["Tell me a joke!"]], inputs=prompt) | |
| gr.Markdown(ARTICLE) | |
| if __name__ == "__main__": | |
| demo.launch() |