Update app.py

app.py

@@ -2,89 +2,122 @@ import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import logging
 import torch
+import os
+from functools import lru_cache
 
 # Logging Setup
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# Model and tokenizer initialization
-model_name = "gpt2"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+# Cache the model and tokenizer loading
+@lru_cache(maxsize=None)
+def load_model_and_tokenizer(model_name):
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(model_name)
+        tokenizer.pad_token = tokenizer.eos_token
+        return model, tokenizer
+    except Exception as e:
+        logging.error(f"Error loading model and tokenizer: {e}")
+        return None, None
+
+# Initialize session state
+if 'generated_codes' not in st.session_state:
+    st.session_state.generated_codes = []
 
-# Set padding token
-tokenizer.pad_token = tokenizer.eos_token
+@st.cache_data
+def generate_code(task_description, max_length, temperature, num_return_sequences, model_name):
+    model, tokenizer = load_model_and_tokenizer(model_name)
+    if model is None or tokenizer is None:
+        return ["Error: Failed to load model and tokenizer."]
 
-def generate_code(task_description, max_length, temperature, num_return_sequences):
     try:
         logging.info(f"Generating code with input: {task_description}")
         prompt = f"Develop code for the following task: {task_description}"
         
         inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
-        
-        # Calculate the maximum new tokens
         max_new_tokens = max_length - inputs.input_ids.shape[1]
         
-        # Generate the output
-        output = model.generate(
-            inputs.input_ids,
-            max_new_tokens=max_new_tokens,  # Use max_new_tokens instead of max_length
-            num_return_sequences=num_return_sequences,
-            temperature=temperature,
-            do_sample=True,
-            no_repeat_ngram_size=2,
-            top_k=50,
-            top_p=0.95,
-            pad_token_id=tokenizer.eos_token_id,  # Explicitly set pad_token_id
-            attention_mask=inputs.attention_mask,  # Pass the attention mask
-        )
+        with torch.no_grad():
+            output = model.generate(
+                inputs.input_ids,
+                max_new_tokens=max_new_tokens,
+                num_return_sequences=num_return_sequences,
+                temperature=temperature,
+                do_sample=True,
+                no_repeat_ngram_size=2,
+                top_k=50,
+                top_p=0.95,
+                pad_token_id=tokenizer.eos_token_id,
+                attention_mask=inputs.attention_mask,
+            )
         
-        # Decode the output
         codes = [tokenizer.decode(seq, skip_special_tokens=True) for seq in output]
-        
         logging.info("Code generation completed successfully.")
         return codes
     except Exception as e:
         logging.error(f"Error generating code: {e}")
         return [f"Error generating code: {e}"]
 
+def save_code(code, file_name):
+    try:
+        with open(file_name, "w") as file:
+            file.write(code)
+        return True
+    except Exception as e:
+        logging.error(f"Error saving code: {e}")
+        return False
+
 def main():
     st.set_page_config(page_title="Advanced Code Generator", layout="wide")
 
     st.title("Advanced Code Generator")
     st.markdown("This application generates code based on the given task description using a text-generation model.")
 
+    # Model Selection
+    model_name = st.selectbox("Select Model", ["gpt2", "gpt2-medium", "gpt2-large"], help="Choose the model for code generation.")
+
     # Input Section
     st.header("Task Description")
     task_description = st.text_area("Describe the task for which you need code:", height=150)
 
     # Options Section
     st.header("Options")
-    max_length = st.slider("Max Length", min_value=50, max_value=1000, value=250, step=50, help="Maximum length of the generated code.")
-    temperature = st.slider("Temperature", min_value=0.1, max_value=1.0, value=0.7, step=0.1, help="Controls the creativity of the generated code.")
-    num_return_sequences = st.slider("Number of Sequences", min_value=1, max_value=5, value=1, step=1, help="Number of code snippets to generate.")
+    col1, col2, col3 = st.columns(3)
+    with col1:
+        max_length = st.slider("Max Length", min_value=50, max_value=1000, value=250, step=50, help="Maximum length of the generated code.")
+    with col2:
+        temperature = st.slider("Temperature", min_value=0.1, max_value=1.0, value=0.7, step=0.1, help="Controls the creativity of the generated code.")
+    with col3:
+        num_return_sequences = st.slider("Number of Sequences", min_value=1, max_value=5, value=1, step=1, help="Number of code snippets to generate.")
 
     # Generate Code Button
     if st.button("Generate Code"):
         if task_description.strip():
             with st.spinner("Generating code..."):
-                generated_codes = generate_code(task_description, max_length, temperature, num_return_sequences)
+                st.session_state.generated_codes = generate_code(task_description, max_length, temperature, num_return_sequences, model_name)
             st.header("Generated Code")
-            for idx, code in enumerate(generated_codes):
-                st.subheader(f"Generated Code {idx + 1}")
-                st.code(code, language='python')
+            for idx, code in enumerate(st.session_state.generated_codes):
+                with st.expander(f"Generated Code {idx + 1}", expanded=True):
+                    st.code(code, language='python')
         else:
             st.error("Please enter a task description.")
 
     # Save Code Section
-    if 'generated_codes' in locals() and generated_codes:
+    if st.session_state.generated_codes:
         st.header("Save Code")
-        selected_code_idx = st.selectbox("Select which code to save:", range(1, len(generated_codes) + 1)) - 1
-        file_name = st.text_input("Enter file name to save:", value="generated_code.py")
-        if st.button("Save", key="save_code"):
+        selected_code_idx = st.selectbox("Select which code to save:", range(1, len(st.session_state.generated_codes) + 1)) - 1
+        col1, col2 = st.columns(2)
+        with col1:
+            file_name = st.text_input("Enter file name to save:", value="generated_code.py")
+        with col2:
+            save_button = st.button("Save", key="save_code")
+        
+        if save_button:
             if file_name:
-                with open(file_name, "w") as file:
-                    file.write(generated_codes[selected_code_idx])
-                st.success(f"Code saved to {file_name}")
+                if save_code(st.session_state.generated_codes[selected_code_idx], file_name):
+                    st.success(f"Code saved to {file_name}")
+                else:
+                    st.error("Failed to save the code. Please try again.")
             else:
                 st.error("Please enter a valid file name.")