Update app.py

app.py

@@ -1,56 +1,37 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import pipeline
 import logging
 import torch
-from functools import lru_cache
 
 # Logging Setup
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# 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 = []
 
+@st.cache_resource
+def get_model_pipeline(model_name):
+    try:
+        code_pipeline = pipeline("text-generation", model=model_name, device=0 if torch.cuda.is_available() else -1)
+        return code_pipeline
+    except Exception as e:
+        logging.error(f"Error loading model pipeline: {e}")
+        return None
+
 @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."]
-    
+    code_pipeline = get_model_pipeline(model_name)
+    if code_pipeline is None:
+        return ["Error: Failed to load model pipeline."]
+
     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)
-        max_new_tokens = max(max_length - inputs.input_ids.shape[1], 1)  # Ensure max_new_tokens is at least 1
-        
-        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,
-            )
+        outputs = code_pipeline(prompt, max_length=max_length, num_return_sequences=num_return_sequences, temperature=temperature)
+        codes = [output['generated_text'] for output in outputs]
         
-        codes = [tokenizer.decode(seq, skip_special_tokens=True) for seq in output]
         logging.info("Code generation completed successfully.")
         return codes
     except Exception as e:
@@ -73,11 +54,13 @@ def main():
     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", 
-        ["EleutherAI/gpt-neo-2.7B", "EleutherAI/gpt-j-6B"], 
-        help="Choose the model for code generation."
-    )
+    model_name = st.selectbox("Select Model", [
+        "EleutherAI/gpt-neo-2.7B", 
+        "EleutherAI/gpt-j-6B", 
+        "bigscience/bloom-1b7",
+        "openai-gpt",
+        "gpt2-xl"
+    ], help="Choose the model for code generation.")
 
     # Input Section
     st.header("Task Description")
@@ -87,7 +70,7 @@ def main():
     st.header("Options")
     col1, col2, col3 = st.columns(3)
     with col1:
-        max_length = st.slider("Max Length", min_value=50, max_value=2048, value=2000, step=25, help="Maximum length of the generated code.")
+        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: