Spaces:
Build error
Build error
File size: 2,466 Bytes
05f1d97 6fe40f1 05f1d97 6fe40f1 05f1d97 6fe40f1 05f1d97 6fe40f1 05f1d97 6fe40f1 05f1d97 6fe40f1 05f1d97 6fe40f1 05f1d97 6fe40f1 05f1d97 6fe40f1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
# Load model and tokenizer
model_name = "microsoft/DialoGPT-small"
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(model_name)
# Move model to GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
def respond(message, chat_history, chat_history_ids):
if not message.strip():
return "", chat_history or [], chat_history_ids, "Please enter a message."
if chat_history is None:
chat_history = []
new_input_ids = tokenizer.encode(message + tokenizer.eos_token, return_tensors="pt").to(device)
input_ids = torch.cat([chat_history_ids, new_input_ids], dim=-1) if chat_history_ids is not None else new_input_ids
try:
chat_history_ids = model.generate(
input_ids,
max_length=200,
pad_token_id=tokenizer.eos_token_id,
no_repeat_ngram_size=3,
do_sample=True,
top_k=50,
top_p=0.95,
temperature=0.8
)
response = tokenizer.decode(
chat_history_ids[:, input_ids.shape[-1]:][0],
skip_special_tokens=True
)
chat_history.append((message, response))
if len(chat_history) > 10:
chat_history = chat_history[-10:]
history_text = "".join([msg + resp + tokenizer.eos_token for msg, resp in chat_history])
chat_history_ids = tokenizer.encode(history_text, return_tensors="pt").to(device)
return "", chat_history, chat_history_ids, None
except Exception as e:
return "", chat_history, chat_history_ids, f"Error: {str(e)}"
def clear_history():
return [], None, None
with gr.Blocks() as demo:
state = gr.State()
gr.Markdown("## DialoGPT Chatbot")
chatbot = gr.Chatbot()
msg = gr.Textbox(label="Your Message", placeholder="Type your message here...")
clear = gr.Button("Clear History")
error = gr.Textbox(label="Error", interactive=False, visible=False)
msg.submit(
respond,
inputs=[msg, chatbot, state],
outputs=[msg, chatbot, state, error]
)
clear.click(
fn=clear_history,
inputs=None,
outputs=[chatbot, state, error],
queue=False
) |