Update README.md

README.md

@@ -42,7 +42,8 @@ This model is not suitable for all use cases due to its limited training time on
 ```python
 import torch
 from transformers import GPT2Tokenizer, GPT2LMHeadModel
-
+start_token = "<|ASSISTANT|>"
+end_token = "<|"
 tokenizer = GPT2Tokenizer.from_pretrained('gpt2-medium')
 model = GPT2LMHeadModel.from_pretrained('gpt2-medium')
 tokenizer.add_special_tokens({'pad_token': '[PAD]'})
@@ -68,9 +69,6 @@ def generate_text(model, tokenizer, prompt, max_length=256):
                             eos_token_id=tokenizer.eos_token_id,
                             attention_mask=attention_mask)
     output_ids = tokenizer.decode(output[0], skip_special_tokens=False)
-    assistant_token_index = output_ids.index('<|ASSISTANT|>') + len('<|ASSISTANT|>')
-    next_token_index = output_ids.find('<|', assistant_token_index)
-    output_ids = output_ids[assistant_token_index:next_token_index]
     return output_ids
 # Loop to interact with the model
 while True:
@@ -78,7 +76,9 @@ while True:
     if prompt == "q":
         break
     output_text = generate_text(model, tokenizer, prompt)
-    print(output_text)
+    text_between_tokens = output_text[output_text.find(start_token) + len(start_token):]
+    out = text_between_tokens[:text_between_tokens.find(end_token)]
+    print(out)
 ```
 ## Deploying and training the model
 The model has been fine-tuned on a specific input format that goes like this ```"<|USER|> {user prompt} <|ASSISTANT|> {model prediction} <|End|>".``` For the best performance from the model the input text should be as follows ```<|USER|> {dataset prompt} <|ASSISTANT|> ``` and the target/label should be as follows ```<|USER|> {dataset prompt} <|ASSISTANT|> {dataset output} <|End|>```