Update app.py

app.py

@@ -1,42 +1,81 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-# Load model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained("namannn/llama2-13b-hyperbolic-cluster-pruned")
-model = AutoModelForCausalLM.from_pretrained("namannn/llama2-13b-hyperbolic-cluster-pruned")
-
-# Streamlit UI components
-st.title("Text Generation with LLaMa2-13b Hyperbolic Model")
-st.write("Enter a prompt below and the model will generate text.")
-
-# User input for prompt
-prompt = st.text_area("Input Prompt", "Once upon a time, in a land far away")
-
-# Slider for controlling the length of the output
-max_length = st.slider("Max Length of Generated Text", min_value=50, max_value=200, value=100)
-
-# Button to trigger text generation
-if st.button("Generate Text"):
-    if prompt:
-        # Encode the prompt text
-        inputs = tokenizer(prompt, return_tensors="pt")
-
-        # Generate text with the model
-        outputs = model.generate(
-            inputs["input_ids"],
-            max_length=max_length,
-            num_return_sequences=1,
-            no_repeat_ngram_size=2,  # You can tune this for diversity
-            do_sample=True,  # Use sampling for diverse generation
-            top_k=50,  # Top-k sampling for diversity
-            top_p=0.95,  # Top-p (nucleus) sampling
-            temperature=0.7  # Control randomness (lower = more deterministic)
-        )
-
-        # Decode and display generated text
-        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        st.subheader("Generated Text:")
-        st.write(generated_text)
-    else:
-        st.warning("Please enter a prompt to generate text.")
+@st.cache_resource
+def load_model_and_tokenizer():
+    """
+    Load model and tokenizer with Streamlit's caching to prevent reloading.
+    @st.cache_resource ensures the model is loaded only once per session.
+    """
+    tokenizer = AutoTokenizer.from_pretrained("namannn/llama2-13b-hyperbolic-cluster-pruned")
+    model = AutoModelForCausalLM.from_pretrained(
+        "namannn/llama2-13b-hyperbolic-cluster-pruned", 
+        # Optional: specify device and precision to optimize loading
+        device_map="auto",  # Automatically distribute model across available GPUs/CPU
+        torch_dtype=torch.float16,  # Use half precision to reduce memory usage
+        low_cpu_mem_usage=True  # Optimize memory usage during model loading
+    )
+    return tokenizer, model
 
+def generate_text(prompt, tokenizer, model, max_length):
+    """
+    Generate text using the loaded model and tokenizer.
+    """
+    # Encode the prompt text
+    inputs = tokenizer(prompt, return_tensors="pt")
+
+    # Generate text with the model
+    outputs = model.generate(
+        inputs["input_ids"], 
+        max_length=max_length, 
+        num_return_sequences=1, 
+        no_repeat_ngram_size=2,
+        do_sample=True, 
+        top_k=50, 
+        top_p=0.95, 
+        temperature=0.7
+    )
+
+    # Decode and return generated text
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return generated_text
+
+def main():
+    # Set page title and icon
+    st.set_page_config(page_title="LLaMa2 Text Generation", page_icon="✍️")
+
+    # Page title and description
+    st.title("Text Generation with LLaMa2-13b Hyperbolic Model")
+    st.write("Enter a prompt below and the model will generate text.")
+
+    # Load model and tokenizer (only once)
+    try:
+        tokenizer, model = load_model_and_tokenizer()
+    except Exception as e:
+        st.error(f"Error loading model: {e}")
+        return
+
+    # User input for prompt
+    prompt = st.text_area("Input Prompt", "Once upon a time, in a land far away")
+
+    # Slider for controlling the length of the output
+    max_length = st.slider("Max Length of Generated Text", min_value=50, max_value=200, value=100)
+
+    # Button to trigger text generation
+    if st.button("Generate Text"):
+        if prompt:
+            try:
+                # Generate text
+                generated_text = generate_text(prompt, tokenizer, model, max_length)
+                
+                # Display generated text
+                st.subheader("Generated Text:")
+                st.write(generated_text)
+            except Exception as e:
+                st.error(f"Error generating text: {e}")
+        else:
+            st.warning("Please enter a prompt to generate text.")
+
+if __name__ == "__main__":
+    main()