Update app.py

app.py

@@ -4,7 +4,7 @@ from transformers import pipeline
 import torch
 
 
-@st.cache(allow_output_mutation=True)
+@st.cache_resource(allow_output_mutation=True)
 def get_model():
     # load base LLM model and tokenizer
 
@@ -21,15 +21,8 @@ def get_model():
 
 tokenizer, model = get_model()
 
-# model_id = "niclasfw/schlager-bot-004"
-
-# model = AutoModelForCausalLM.from_pretrained(model_id)
-# tokenizer = AutoTokenizer.from_pretrained(model_id)
-
-# generator = pipeline(task="text-generation", model=model_id, tokenizer=model_id)
-
 st.title('Schlager Bot')
-user_input = st.text_area('Enter verse (minimum of 15 words): ')
+user_input = st.text_area('Enter verse: ')
 button = st.button('Generate Lyrics')
 
 
@@ -42,14 +35,9 @@ if user_input and button:
 
     ### Response:
     """
-    # output = generator(prompt, do_sample=True, max_new_tokens=500, top_p=0.75, temperature=0.95, top_k=15)
-    # st.write("Prompt: ", user_input)
-    # input = tokenizer(prompt, padding=True, return_tensors="pt")
-    # generate_ids = model.generate(input.input_ids, max_length=500, top_p=0.75, temperature=0.95, top_k=15)
-    # output = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
     input_ids = tokenizer(prompt, return_tensors="pt", truncation=True).input_ids.cuda()
     outputs = model.generate(input_ids=input_ids, pad_token_id=tokenizer.eos_token_id, max_new_tokens=500, do_sample=True, top_p=0.75, temperature=0.95, top_k=15)
 
-    st.write(output)
+    st.write(outputs)