refactor

Dockerfile

@@ -20,7 +20,6 @@ RUN pip install -r requirements.txt
 
 EXPOSE 7860
 
-RUN mkdir -p /home/user/app/.secrets
 RUN --mount=type=secret,id=HF_TOKEN,mode=0444,required=true
 
 CMD ["python", "app.py"]

app.py

@@ -9,49 +9,55 @@ 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)
+# 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])
 
-# 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,
-# )
+#     # 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)
 
-# def llama2(input_text):
-#     sequences = pipeline(
-#         input_text,
-#         do_sample=True,
-#         top_k=10,
-#         num_return_sequences=1,
-#         eos_token_id=tokenizer.eos_token_id,
-#         max_length=200,
-#     )
-#     output_text = ""
-#     for seq in sequences:
-#         output_text += seq["generated_text"] + "\n"
-#     return output_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 "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=google_flan, 
+    fn=llama2, 
     inputs="text", 
+    # outputs="text", 
     outputs=gr.Textbox(
             label="LLM",
             lines=5,