Commit
•
55c8e0a
1
Parent(s):
66f3bf5
refactor
Browse files- app.py +3 -1
- docker-compose.yml +1 -2
app.py
CHANGED
@@ -12,11 +12,13 @@ else:
|
|
12 |
logging.info("Running on CPU")
|
13 |
|
14 |
if "googleflan" == os.environ.get("MODEL"):
|
|
|
15 |
model = "google/flan-t5-small"
|
16 |
pipeline = transformers.pipeline("text2text-generation", model=model)
|
17 |
def model_func(input_text, request: gr.Request):
|
18 |
return pipeline(input_text)
|
19 |
elif "llama" == os.environ.get("MODEL"):
|
|
|
20 |
model = "meta-llama/Llama-2-7b-chat-hf"
|
21 |
tokenizer = AutoTokenizer.from_pretrained(
|
22 |
model,
|
@@ -50,7 +52,7 @@ elif "summary" == os.environ.get("MODEL"):
|
|
50 |
model="facebook/bart-large-cnn"
|
51 |
summarizer = transformers.pipeline("summarization", model=model)
|
52 |
def model_func(input_text, request: gr.Request):
|
53 |
-
return summarizer(input_text, max_length=
|
54 |
|
55 |
input_label = "How can I help?"
|
56 |
if "summary" == os.environ.get("MODEL"):
|
|
|
12 |
logging.info("Running on CPU")
|
13 |
|
14 |
if "googleflan" == os.environ.get("MODEL"):
|
15 |
+
### Fast/small model used to debug UI on local machine
|
16 |
model = "google/flan-t5-small"
|
17 |
pipeline = transformers.pipeline("text2text-generation", model=model)
|
18 |
def model_func(input_text, request: gr.Request):
|
19 |
return pipeline(input_text)
|
20 |
elif "llama" == os.environ.get("MODEL"):
|
21 |
+
### Works on CPU but runtime is > 4 minutes
|
22 |
model = "meta-llama/Llama-2-7b-chat-hf"
|
23 |
tokenizer = AutoTokenizer.from_pretrained(
|
24 |
model,
|
|
|
52 |
model="facebook/bart-large-cnn"
|
53 |
summarizer = transformers.pipeline("summarization", model=model)
|
54 |
def model_func(input_text, request: gr.Request):
|
55 |
+
return summarizer(input_text, max_length=300, min_length=30, do_sample=False)[0]["summary_text"]
|
56 |
|
57 |
input_label = "How can I help?"
|
58 |
if "summary" == os.environ.get("MODEL"):
|
docker-compose.yml
CHANGED
@@ -11,8 +11,7 @@ services:
|
|
11 |
working_dir: /home/user/app
|
12 |
environment:
|
13 |
- HF_TOKEN=${HF_TOKEN}
|
14 |
-
- MODEL=
|
15 |
-
# - MODEL=googleflan
|
16 |
stdin_open: true
|
17 |
tty: true
|
18 |
restart: always
|
|
|
11 |
working_dir: /home/user/app
|
12 |
environment:
|
13 |
- HF_TOKEN=${HF_TOKEN}
|
14 |
+
- MODEL=googleflan
|
|
|
15 |
stdin_open: true
|
16 |
tty: true
|
17 |
restart: always
|