Upload app.py

app.py

@@ -9,7 +9,7 @@ def main():
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     tokenizer, EMPTY_TOKENS = load_tokenizer(device)
 
-    st.title("TinyStories Transformer LLM Demo")
+    st.title("TinyStories Transformer Demo 🤖")
 
     st.subheader("Data and Training")
     
@@ -24,7 +24,7 @@ def main():
     st.subheader("How Do I Prompt?")
     st.markdown(
         """
-        Instead of generating a new story from scratch, you can "prompt" the model by writing the first few words or sentences of a story, and let it finish from there. It can even jump in mid-sentence!
+        Instead of generating a new story from scratch, you can "prompt" the model by writing the first few words of a story, and let it finish from there. It can even jump in mid-sentence!
 
         The model can struggle with some prompts, especially those outside of its limited domain. If a response isn't satisfactory, try repeating the generation, or make the following modifications:
         """
@@ -73,8 +73,8 @@ def main():
     # model_version = st.radio("Which model would you like to use?", ["smoll", "beeg"])
     # small_model = load_casey_model(tokenizer, device)
     model = load_big_model(tokenizer, device)
-    #model.to('cuda')
-    #model.cuda()
+    model.to('cuda')
+    model.cuda()
 
 
     if st.button('Write my story!'):