Llama2 init

Dockerfile

@@ -6,15 +6,11 @@ RUN useradd -m -u 1000 user
 # Switch to the "user" user
 USER user
 
-# Set home to the user's home directory
-ENV HOME=/home/user \
-    PATH=/home/user/.local/bin:$PATH
-
 # Set the working directory to the user's home directory
-WORKDIR $HOME/app
+WORKDIR /home/user/app
 
 # Copy the current directory contents into the container at $HOME/app setting the owner to the user
-COPY --chown=user . $HOME/app
+COPY --chown=user . /home/user/app
 
 RUN pip install -r requirements.txt
 

README.md

@@ -17,7 +17,7 @@ See live app on [Hugging Face](https://huggingface.co/spaces/AdamNovotnyCom/llam
 
 Start
 
-    docker-compose -f docker-compose.yml up --detach llama2hf
+    export HF_TOKEN=paste_HF_token && docker-compose -f docker-compose.yml up llama2hf
 
 View app in browser
 
@@ -28,4 +28,5 @@ Exec command
     docker exec -it llama2hf bash -c 'pip install torch==2.1.*'
 
 ## References
-- https://huggingface.co/blog/llama2
+- [huggingface.co/llama2](https://huggingface.co/blog/llama2)
+- [demo-docker-gradio](https://huggingface.co/spaces/sayakpaul/demo-docker-gradio/tree/main)

app.py

@@ -1,11 +1,41 @@
 import gradio as gr
-from transformers import pipeline
+import os
+import torch
+import transformers
+from transformers import AutoTokenizer
 
-pipe_flan = pipeline("text2text-generation", model="google/flan-t5-small")
+# pipe_flan = transformers.pipeline("text2text-generation", model="google/flan-t5-small")
+# def google_flan(input_text):
+#     return pipe_flan(input_text)["generated_text"]
 
-def llm(input_text):
-    return pipe_flan(input_text)["generated_text"]
+model = "meta-llama/Llama-2-7b-chat-hf"
+tokenizer = AutoTokenizer.from_pretrained(
+    model,
+    token=os.environ["HF_TOKEN"],
+)
+pipeline = transformers.pipeline(
+    "text-generation",
+    model=model,
+    torch_dtype=torch.float16,
+    device_map="auto",
+    token=os.environ["HF_TOKEN"],
+    low_cpu_mem_usage=False,
+)
 
-demo = gr.Interface(fn=llm, inputs="text", outputs="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
+
+demo = gr.Interface(fn=llama2, inputs="text", outputs="text")
     
 demo.launch(server_name="0.0.0.0", server_port=7860)

docker-compose.yml

@@ -9,6 +9,8 @@ services:
     volumes:
       - ./:/home/user/app
     working_dir: /home/user/app
+    environment:
+      - HF_TOKEN=${HF_TOKEN}
     stdin_open: true
     tty: true
     restart: always

requirements.txt

@@ -1,3 +1,4 @@
+accelerate==0.23.*
 gradio==3.47.*
 pandas==2.1.*
 pytest==7.4.*