refactor

app.py

@@ -23,11 +23,6 @@ elif "summary_bart" == os.environ.get("MODEL"):
     summarizer = transformers.pipeline("summarization", model=model)
     def model_func(input_text, request: gr.Request):
         return summarizer(input_text, max_length=130, min_length=30, do_sample=False)[0]["summary_text"]
-elif "summary_pegasus" == os.environ.get("MODEL"):
-    model="pegasus-cnn_dailymail"
-    summarizer = transformers.pipeline("summarization", model=model)
-    def model_func(input_text, request: gr.Request):
-        return summarizer(input_text, max_length=130, min_length=30, do_sample=False)[0]["summary_text"]
 elif "llama" == os.environ.get("MODEL"):
     ### Works on CPU but runtime is > 4 minutes
     model = "meta-llama/Llama-2-7b-chat-hf"
@@ -62,7 +57,7 @@ elif "llama" == os.environ.get("MODEL"):
 
 # UI: Gradio
 input_label = "How can I help?"
-if "summary" == os.environ.get("MODEL"):
+if "summary" in os.environ.get("MODEL"):
     input_label = "Enter text to summarize"
 demo = gr.Interface(
     fn=model_func, 

docker-compose.yml

@@ -11,7 +11,7 @@ services:
     working_dir: /home/user/app
     environment:
       - HF_TOKEN=${HF_TOKEN}
-      - MODEL=summary_bart
+      - MODEL=googleflan
     stdin_open: true
     tty: true
     restart: always