DialoGPTLarge / app.py
Andreyalth's picture
agregar app
dfa4de0
# Load model directly
from transformers import AutoTokenizer, AutoModelForCausalLM
import gradio as gr
import torch
""" from huggingface_hub import login
login(token="token", add_to_git_credential=True) """
model_name = "microsoft/DialoGPT-large"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
model.generation_config.pad_token_id = model.generation_config.eos_token_id
chat_history_ids = None
def generate_response(user_input, history = []):
global chat_history_ids # Para modificar la variable global
new_user_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt')
if len(history) > 0:
bot_input_ids = torch.cat([chat_history_ids, new_user_input_ids], dim=-1)
else:
bot_input_ids = new_user_input_ids
chat_history_ids = model.generate(bot_input_ids, max_length=10000, pad_token_id=tokenizer.eos_token_id)
response = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
# Para depuración
#print("user", new_user_input_ids)
#print("bot",chat_history_ids[:, bot_input_ids.shape[-1]:][0])
#print(chat_history_ids)
return response
chatbot_interface = gr.ChatInterface(fn=generate_response)
demo = gr.TabbedInterface([
chatbot_interface
], ["chatbot"])
demo.launch()