Update app.py

app.py

@@ -1,18 +1,26 @@
 import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Load the model and tokenizer
+# Load the model and tokenizer from the Hugging Face repository
 model_name = "aifeifei798/DarkIdol-Llama-3.1-8B-Instruct-1.2-Uncensored"
+st.write(f"Loading model {model_name}...")
+
+# The model and tokenizer will be downloaded and cached locally
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
-# Define the Streamlit app
-st.title("Text Generation with DarkIdol-Llama")
-prompt = st.text_input("Enter your prompt:")
-
-if st.button("Generate"):
+# Function to generate text based on user input
+def generate_text(prompt):
     inputs = tokenizer(prompt, return_tensors="pt")
     outputs = model.generate(**inputs, max_new_tokens=100)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    st.write("Response:")
-    st.write(response)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+# Streamlit UI
+st.title("DarkIdol Text Generation")
+user_input = st.text_area("Enter your prompt:")
+
+if st.button("Generate"):
+    with st.spinner("Generating..."):
+        output_text = generate_text(user_input)
+    st.write("Generated Text:")
+    st.write(output_text)