Pushing updates to huggingface space.

app.py

@@ -1,38 +1,42 @@
 import streamlit as st
 import torch
-from transformers import AutoModelForCausalLM
-from transformers import AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-model_id = "niclasfw/schlager-bot-004"
+@st.cache(allow_output_mutation=True)
+def get_model():
+    # load base LLM model and tokenizer
 
-# load base LLM model and tokenizer
-model = AutoModelForCausalLM.from_pretrained(
+    model_id = "niclasfw/schlager-bot-004"
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model = AutoModelForCausalLM.from_pretrained(
     model_id,
     low_cpu_mem_usage=True,
     torch_dtype=torch.float16,
     load_in_4bit=True,
-)
-tokenizer = AutoTokenizer.from_pretrained(model_id)
+    )
 
-st.title('Schlager Bot')
+    return tokenizer, model
+
+tokenizer, model = get_model()
 
-print("Enter verse (minimum of 15 words): ")
-sample = input()
+st.title('Schlager Bot')
+user_input = st.text_area('Enter verse (minimum of 15 words): ')
+button = st.button('Generate Lyrics')
 
-prompt = f"""### Instruction:
-Benuzte den gegebenen Input um ein Schlager Lied zu schreiben.
 
-### Input:
-{sample}
+if user_input and button:
+    prompt = f"""### Instruction:
+    Benuzte den gegebenen Input um ein Schlager Lied zu schreiben.
 
-### Response:
-"""
+    ### Input:
+    {user_input}
 
-query = st.text_input(prompt, key="input")
+    ### Response:
+    """
+    input_ids = tokenizer(prompt, return_tensors="pt", truncation=True).input_ids
+    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("Prompt: ", user_input)
+    st.write("**************")
+    st.write(outputs)
 
-input_ids = tokenizer(prompt, return_tensors="pt", truncation=True).input_ids
-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)
 
-print(f"Prompt:\n{sample}\n")
-print("**************")
-st.write(outputs)

requirements.txt

@@ -0,0 +1,3 @@
+streamlit
+torch
+transformers