Update app.py

app.py

@@ -1,31 +1,30 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-# Load the model and tokenizer from Hugging Face
-model_name = "Salesforce/codet5-small"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+@st.cache_resource
+def load_model():
+    model_name = "Salesforce/codet5-small"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+    return tokenizer, model
 
-# Streamlit UI
-st.title("Code Generator")
-st.write("Generate code snippets from natural language prompts using CodeT5!")
+# Load the model and tokenizer (cached)
+with st.spinner("Loading model..."):
+    tokenizer, model = load_model()
 
-# Input for natural language prompt
-prompt = st.text_area("Enter your coding task:", placeholder="Write a Python function to calculate the factorial of a number.")
+# Streamlit UI
+st.title("Code Generator with Hugging Face")
+st.write("Generate code snippets from natural language prompts!")
 
-# Slider to control output length
-max_length = st.slider("Maximum length of generated code:", 20, 200, 50)
+prompt = st.text_area("Enter your coding task:", placeholder="Write a Python function to calculate factorial.")
+max_length = st.slider("Select maximum length of generated code:", min_value=20, max_value=200, value=50, step=10)
 
-# Button to trigger code generation
 if st.button("Generate Code"):
     if prompt.strip():
-        # Tokenize and generate code
-        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True)
-        outputs = model.generate(inputs.input_ids, max_length=max_length, num_beams=4, early_stopping=True)
-        generated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-        # Display generated code
-        st.write("### Generated Code:")
-        st.code(generated_code, language="python")
+        with st.spinner("Generating code..."):
+            inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True)
+            outputs = model.generate(inputs.input_ids, max_length=max_length, num_beams=4, early_stopping=True)
+            generated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
+            st.text_area("Generated Code:", generated_code, height=200)
     else:
-        st.warning("Please enter a prompt to generate code.")
+        st.warning("Please enter a prompt!")