Update app.py

app.py

@@ -9,6 +9,8 @@ MODEL_ID = "LMSeed/GPT2-small-distilled-900M_None_ppo-1000K-seed42"
 device = 0 if torch.cuda.is_available() else -1
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token_id
 model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
 
 if torch.cuda.is_available():
@@ -38,18 +40,18 @@ def generate_reply(prompt, max_new_tokens, temperature, top_p):
 def clean_reply(text):
     
     text = text.strip()
-    
-    for prefix in ["Assistant:", "assistant:", "User:", "user:"]:
-        if text.startswith(prefix):
-            text = text[len(prefix):].strip()
+    return text
+    # for prefix in ["Assistant:", "assistant:", "User:", "user:"]:
+    #     if text.startswith(prefix):
+    #         text = text[len(prefix):].strip()
             
-    lines = [l.strip() for l in text.split("\n")]
-    lines = [l for l in lines if l]
+    # lines = [l.strip() for l in text.split("\n")]
+    # lines = [l for l in lines if l]
 
-    if len(lines) == 0:
-        return ""
+    # if len(lines) == 0:
+    #     return ""
 
-    return lines[0] 
+    # return lines[0] 
     
 def chat_with_model(user_message, chat_history, max_new_tokens=256, temperature=0.8, top_p=0.9):
 
@@ -58,8 +60,8 @@ def chat_with_model(user_message, chat_history, max_new_tokens=256, temperature=
 
     # Build conversation history
     # history_text = "The following is a friendly conversation between a human and an AI assistant.\n"
-    history_text = "The following is a friendly conversation between a human and an AI story-telling assistant. \
-            The assistant should tell a story according to human's requirment.\n"
+    history_text = "" #"The following is a friendly conversation between a human and an AI story-telling assistant. \
+            # The assistant should tell a story according to human's requirment.\n"
 
     for msg in chat_history:
         role = "Human" if msg["role"] == "user" else "AI"