Update blueprints/summarize.py

blueprints/summarize.py

@@ -14,7 +14,7 @@ MODEL_CACHE = {
     "pipeline": None
 }
 
-def get_pipeline(model_name, task_type):
+def get_pipeline(model_name, task_type, hf_token):
     """
     Retrieves a pipeline from cache or loads it if it's new.
     """
@@ -25,7 +25,6 @@ def get_pipeline(model_name, task_type):
         return MODEL_CACHE["pipeline"]
 
     # Authentication
-    hf_token = os.getenv("HF_TOKEN")
     if hf_token: 
         HfFolder.save_token(hf_token)
 
@@ -52,7 +51,7 @@ def run_inference_logic(config):
     model_id = config['model_name']
     text = config['text']
     task_type = config['task_type']
-
+    hf_token = config['hf_token']
     pipe = get_pipeline(model_id, task_type)
 
     if task_type == 'TOKEN_CLS':
@@ -93,6 +92,7 @@ def api_summarize():
     config = {
         "text": data['text'],
         "model_name": data.get('model_name', "facebook/bart-large-cnn"),
+        "hf_token": data.get('hf_token',os.getenv("HF_TOKEN"),'')
         # We force this for the specific summarization UI, 
         # but the backend logic supports others.
         "task_type": "SEQ_2_SEQ_LM"