new cosmetics

app.py

@@ -46,16 +46,25 @@ st.title("Rap Verse Generation V1 Demo")
 st.header("Supported Artists")
 st.write("Lupe Fiasco, Common, Jay-Z, Yasiin Bey, Ab-Soul, Rakim")
 
+prompt_placeholder = st.empty()
 display_placeholder = st.empty()
-display_placeholder.write("Ask Rap-Mistral Something")
+
+prompt_placeholder.write("Ask Rap-Mistral Something")
+display_placeholder.write("")
 
 question = st.chat_input("Write a verse in the style of Lupe Fiasco")
 if question:
     display_placeholder.write("Loading...")
+    prompt_placeholder.write(question)
     prompt = build_prompt(question)
     inputs = tokenizer(prompt, return_tensors="pt")
     model_inputs = inputs.to('cuda')
     generated_ids = model.generate(**model_inputs, max_new_tokens=300, do_sample=True, pad_token_id=tokenizer.eos_token_id)
-    response = tokenizer.batch_decode(generated_ids)
-    display_placeholder.write(response[0])
+    response = tokenizer.batch_decode(generated_ids)[0]
+    end_of_inst = response.find("[/INST]") + len("[/INST]")
+    if end_of_inst > -1:
+        actual_response = response[end_of_inst:].strip()
+    else:
+        actual_response = response
+    display_placeholder.write(actual_response)