Update app.py

app.py

@@ -1,6 +1,13 @@
 import streamlit as st
 
 from transformers import GPT2Tokenizer, GPT2LMHeadModel
+model_name = 'KhantKyaw/Chat_GPT-2'
+tokenizer = GPT2Tokenizer.from_pretrained(model_name)
+model = GPT2LMHeadModel.from_pretrained(model_name)
+
+special_tokens_dict = {'bos_token': '<BOS>', 'eos_token': '<EOS>', 'sep_token': '<SEP>', 'pad_token': '<PAD>'}
+tokenizer.add_special_tokens(special_tokens_dict)
+model.resize_token_embeddings(len(tokenizer))
 
 # Function to generate a response
 def generate_response(input_text):
@@ -20,6 +27,7 @@ def generate_response(input_text):
         top_p=0.95,
         no_repeat_ngram_size=2,
         pad_token_id=tokenizer.eos_token_id,
+        eos_token_id=tokenizer.eos_token_id,
         #early_stopping=True,
         do_sample=True
     )
@@ -30,20 +38,14 @@ def generate_response(input_text):
     # Extract the generated response after the [Bot] marker
     bot_response_start = full_generated_text.find('[Bot]') + len('[Bot]')
     bot_response = full_generated_text[bot_response_start:]
-    
-    # Trim the response to end at the last period within the specified max_length
-    last_period_index = bot_response.rfind('.')
-    if last_period_index != -1:
-        bot_response = bot_response[:last_period_index + 1]
-
-    return bot_response.strip()
+    return bot_response
 
 # Load pre-trained model tokenizer (vocabulary) and model
 model_name = 'KhantKyaw/Chat_GPT-2'
 tokenizer = GPT2Tokenizer.from_pretrained(model_name)
 model = GPT2LMHeadModel.from_pretrained(model_name)
 
-st.title("Echo Bot")
+st.title("Chat_GPT-2 Bot")
 
 # Initialize chat history
 if "messages" not in st.session_state: