Update app.py

app.py

@@ -15,7 +15,6 @@ tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
 # Load the Lora model
 model = PeftModel.from_pretrained(model, peft_model_id)
 
-
 def make_inference(synopsis):
     batch = tokenizer(
         f"Below is a one-sentence synopsis, please write a captivating short story based on this synopsis.\n\n### Synopsis:\n{synopsis}\n\n### Short Story:\n", return_tensors='pt',
@@ -24,7 +23,11 @@ def make_inference(synopsis):
     with torch.cuda.amp.autocast():
         output_tokens = model.generate(**batch, max_new_tokens=200)
 
-    return tokenizer.decode(output_tokens[0], skip_special_tokens=True)
+    full_output = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
+    short_story = full_output.split("### Short Story:\n")[-1].strip()
+
+    return short_story
+
 
 
 if __name__ == "__main__":