Update app.py

app.py

@@ -15,11 +15,12 @@ model.to('cpu')
 
 # Define the function that generates text from a prompt
 def generate_text(prompt, temperature, top_p):
-    input_tokens = tokenizer.encode(prompt, return_tensors='pt')
+    prompt_with_eos = " #CONTEXT# " + prompt + " #TOPIC# "  # Add the string "EOS" to the end of the prompt
+    input_tokens = tokenizer.encode(prompt_with_eos, return_tensors='pt')
     input_tokens = input_tokens.to('cpu')
     
-    generated_text = prompt  # Start with the initial prompt
-    prompt_length = len(generated_text)  # Capture the length of the initial prompt
+    generated_text = prompt_with_eos  # Start with the initial prompt plus "EOS"
+    prompt_length = len(generated_text)  # Capture the length of the initial prompt plus "EOS"
 
     for _ in range(80):  # Adjust the range to control the number of tokens generated
         with torch.no_grad():
@@ -40,7 +41,7 @@ def generate_text(prompt, temperature, top_p):
         generated_text += decoded_token  # Append the new token to the generated text
         if decoded_token == "#":  # Stop if the end of sequence token is generated
             break
-        yield generated_text[prompt_length:]  # Yield the generated text excluding the initial prompt
+        yield generated_text[prompt_length:]  # Yield the generated text excluding the initial prompt plus "EOS"
 
 # Create a Gradio interface with a text input, sliders for temperature and top_p, and a text output
 interface = gr.Interface(