Update app.py

app.py

@@ -38,9 +38,9 @@ model.to(device)
 def generate_text(prompt):
     # Prevent gradient calculation to speed up inference
     with torch.no_grad():
-        inputs = tokenizer(prompt, return_tensors="pt").to(device)
+        inputs = tokenizer("<|STARTOFTEXT|> <|USER|> " + prompt + " <|BOT|> ", return_tensors="pt").to(device)
         outputs = model.generate(inputs['input_ids'], max_length=100, num_return_sequences=1, do_sample=False)
-        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=False)
     return generated_text
 
 # Set up the Gradio interface