multiple models

Dockerfile

@@ -20,6 +20,8 @@ RUN pip install -r requirements.txt
 
 EXPOSE 7860
 
+ENV MODEL=llama
+
 RUN --mount=type=secret,id=HF_TOKEN,mode=0444,required=true
 
 CMD ["python", "app.py"]

Dockerfile_dev

@@ -20,5 +20,7 @@ RUN pip install -r requirements.txt
 
 EXPOSE 7860
 
+ENV MODEL=googleflan
+
 # with reload
 CMD ["gradio", "app.py"]

app.py

@@ -7,64 +7,61 @@ from transformers import AutoTokenizer
 
 logging.basicConfig(level=logging.INFO)
 
-print("APP startup")
-
-# pipe_flan = transformers.pipeline("text2text-generation", model="google/flan-t5-small")
-# def google_flan(input_text, request: gr.Request):
-#     print(f"Input request: {input_text}")
-#     print(request.query_params)
-#     print(os.environ.get("HF_TOKEN")[:5])
-#     logging.info(os.environ.get("HF_TOKEN")[:5])
-
-#     # with open("/home/user/app/.secrets/HF_TOKEN", "r") as f:
-#     #     var = f.readlines()[0].strip()
-#     #     print(var[:5])
-#     #     logging.info(var[:5])
-#     return pipe_flan(input_text)
-
-model = "meta-llama/Llama-2-7b-chat-hf"
-tokenizer = AutoTokenizer.from_pretrained(
-    model,
-    token=os.environ.get("HF_TOKEN"),
-)
-pipeline = transformers.pipeline(
-    "text-generation",
-    model=model,
-    torch_dtype=torch.float16,
-    device_map="auto",
-    token=os.environ.get("HF_TOKEN"),
-    low_cpu_mem_usage=True,
-)
-
-def llama2(input_text, request: gr.Request):
-    sequences = pipeline(
-        input_text,
-        do_sample=True,
-        top_k=10,
-        num_return_sequences=1,
-        eos_token_id=tokenizer.eos_token_id,
-        max_length=200,
+if "googleflan" == os.environ.get("MODEL"):
+    model = "google/flan-t5-small"
+    logging.info(f"APP startup. Model {model}")
+    pipe_flan = transformers.pipeline("text2text-generation", model=model)
+    def model_func(input_text, request: gr.Request):
+        print(f"Input request: {input_text}")
+        print(request.query_params)
+        print(os.environ.get("HF_TOKEN")[:5])
+        logging.info(os.environ.get("HF_TOKEN")[:5])
+        return pipe_flan(input_text)
+elif "llama" == os.environ.get("MODEL"):
+    model = "meta-llama/Llama-2-7b-chat-hf"
+    logging.info(f"APP startup. Model {model}")
+    tokenizer = AutoTokenizer.from_pretrained(
+        model,
+        token=os.environ.get("HF_TOKEN"),
+    )
+    pipeline = transformers.pipeline(
+        "text-generation",
+        model=model,
+        torch_dtype=torch.float16,
+        device_map="auto",
+        token=os.environ.get("HF_TOKEN"),
     )
-    if "name" in list(request.query_params):
-        output_text = f"{request.query_params['name']}:\n"
-    else:
-        output_text = ""
-    for seq in sequences:
-        output_text += seq["generated_text"] + "\n"
-    return output_text
+
+    def model_func(input_text, request: gr.Request):
+        sequences = pipeline(
+            input_text,
+            do_sample=True,
+            top_k=10,
+            num_return_sequences=1,
+            eos_token_id=tokenizer.eos_token_id,
+            max_length=200,
+        )
+        if "name" in list(request.query_params):
+            output_text = f"{request.query_params['name']}:\n"
+        else:
+            output_text = ""
+        for seq in sequences:
+            output_text += seq["generated_text"] + "\n"
+        return output_text
 
 demo = gr.Interface(
-    # fn=google_flan, 
-    fn=llama2, 
-    inputs="text", 
-    # outputs="text", 
+    fn=model_func, 
+    inputs=gr.Textbox(
+            label="How can I help?",
+            lines=3,
+            value="",
+    ),
     outputs=gr.Textbox(
             label="LLM",
             lines=5,
             value="",
     ),
     allow_flagging=False, 
-    title="How can I help?",
     theme=gr.themes.Default(primary_hue="blue", secondary_hue="pink")
 )
     

docker-compose.yml

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