Update app.py

app.py

@@ -4,19 +4,21 @@ from transformers import pipeline, Pipeline
 from transformers.pipelines import PipelineException
 from huggingface_hub.utils import ModelNotFoundError
 import logging
+import os
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 # Initialize Hugging Face Hub search component
-search_in = HuggingfaceHubSearch(api_key="hf_YourAPITokenHere", submit_on_select=True)
+HF_TOKEN = os.getenv("HF_TOKEN")
+search_in = HuggingfaceHubSearch(api_key=HF_TOKEN, submit_on_select=True)
 
 # Function to load the selected model and create a pipeline
 def load_model(model_id):
     try:
         logger.info(f"Loading model: {model_id}")
-        model_pipeline = pipeline(model=model_id)
+        model_pipeline = pipeline(model="Salesforce/codet5-small")
         logger.info("Model loaded successfully.")
         return model_pipeline
     except ModelNotFoundError: