| | |
| | import streamlit as st |
| | import torch |
| | import transformers |
| | from transformers import AutoModelForCausalLM, AutoTokenizer |
| |
|
| | @st.cache(hash_funcs= |
| | {transformers.models.gpt2.tokenization_gpt2_fast.GPT2TokenizerFast: hash}, |
| | suppress_st_warning=True, allow_output_mutation=True) |
| |
|
| | def load_data(): |
| | tokenizer = AutoTokenizer.from_pretrained("Rubiksman1006/gpt-neo-2.7b-monika-fp16") |
| | model = AutoModelForCausalLM.from_pretrained("Rubiksman1006/gpt-neo-2.7b-monika-fp16") |
| | return tokenizer, model |
| |
|
| | tokenizer, model = load_data() |
| |
|
| | st.write("Welcome to the Chatbot. I am still learning, please be patient") |
| | input = st.text_input('User:') |
| | if 'count' not in st.session_state or st.session_state.count == 6: |
| | st.session_state.count = 0 |
| | st.session_state.chat_history_ids = None |
| | st.session_state.old_response = '' |
| | else: |
| | st.session_state.count += 1 |
| |
|
| |
|
| | new_user_input_ids = tokenizer.encode(input + tokenizer.eos_token, return_tensors='pt') |
| |
|
| | bot_input_ids = torch.cat([st.session_state.chat_history_ids, new_user_input_ids], dim=-1) if st.session_state.count > 1 else new_user_input_ids |
| |
|
| | st.session_state.chat_history_ids = model.generate(bot_input_ids, max_length=5000, pad_token_id=tokenizer.eos_token_id) |
| |
|
| | response = tokenizer.decode(st.session_state.chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True) |
| |
|
| | if st.session_state.old_response == response: |
| | bot_input_ids = new_user_input_ids |
| | |
| | st.session_state.chat_history_ids = model.generate(bot_input_ids, max_length=5000, pad_token_id=tokenizer.eos_token_id) |
| | response = tokenizer.decode(st.session_state.chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True) |
| |
|
| | st.write(f"Chatbot: {response}") |
| | st.session_state.old_response = response |
| |
|
