Commit
•
fc25e76
1
Parent(s):
65c9805
refactor
Browse files- app.py +1 -6
- docker-compose.yml +1 -1
app.py
CHANGED
@@ -23,11 +23,6 @@ elif "summary_bart" == os.environ.get("MODEL"):
|
|
23 |
summarizer = transformers.pipeline("summarization", model=model)
|
24 |
def model_func(input_text, request: gr.Request):
|
25 |
return summarizer(input_text, max_length=130, min_length=30, do_sample=False)[0]["summary_text"]
|
26 |
-
elif "summary_pegasus" == os.environ.get("MODEL"):
|
27 |
-
model="pegasus-cnn_dailymail"
|
28 |
-
summarizer = transformers.pipeline("summarization", model=model)
|
29 |
-
def model_func(input_text, request: gr.Request):
|
30 |
-
return summarizer(input_text, max_length=130, min_length=30, do_sample=False)[0]["summary_text"]
|
31 |
elif "llama" == os.environ.get("MODEL"):
|
32 |
### Works on CPU but runtime is > 4 minutes
|
33 |
model = "meta-llama/Llama-2-7b-chat-hf"
|
@@ -62,7 +57,7 @@ elif "llama" == os.environ.get("MODEL"):
|
|
62 |
|
63 |
# UI: Gradio
|
64 |
input_label = "How can I help?"
|
65 |
-
if "summary"
|
66 |
input_label = "Enter text to summarize"
|
67 |
demo = gr.Interface(
|
68 |
fn=model_func,
|
|
|
23 |
summarizer = transformers.pipeline("summarization", model=model)
|
24 |
def model_func(input_text, request: gr.Request):
|
25 |
return summarizer(input_text, max_length=130, min_length=30, do_sample=False)[0]["summary_text"]
|
|
|
|
|
|
|
|
|
|
|
26 |
elif "llama" == os.environ.get("MODEL"):
|
27 |
### Works on CPU but runtime is > 4 minutes
|
28 |
model = "meta-llama/Llama-2-7b-chat-hf"
|
|
|
57 |
|
58 |
# UI: Gradio
|
59 |
input_label = "How can I help?"
|
60 |
+
if "summary" in os.environ.get("MODEL"):
|
61 |
input_label = "Enter text to summarize"
|
62 |
demo = gr.Interface(
|
63 |
fn=model_func,
|
docker-compose.yml
CHANGED
@@ -11,7 +11,7 @@ services:
|
|
11 |
working_dir: /home/user/app
|
12 |
environment:
|
13 |
- HF_TOKEN=${HF_TOKEN}
|
14 |
-
- MODEL=
|
15 |
stdin_open: true
|
16 |
tty: true
|
17 |
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
|