Spaces:
Runtime error
Runtime error
Commit
·
36981a1
1
Parent(s):
048132e
Update summarize.py
Browse files- summarize.py +5 -5
summarize.py
CHANGED
|
@@ -21,11 +21,11 @@ def load_model_and_tokenizer(model_name: str) -> tuple:
|
|
| 21 |
:return tuple: a tuple containing the model and tokenizer
|
| 22 |
"""
|
| 23 |
MODEL_OPTIONS = {
|
| 24 |
-
"
|
| 25 |
-
"
|
| 26 |
-
"
|
| 27 |
-
"
|
| 28 |
-
"
|
| 29 |
}
|
| 30 |
selected_model_identifier = MODEL_OPTIONS.get(model_name)
|
| 31 |
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
| 21 |
:return tuple: a tuple containing the model and tokenizer
|
| 22 |
"""
|
| 23 |
MODEL_OPTIONS = {
|
| 24 |
+
"Model_1": "pszemraj/long-t5-tglobal-base-16384-book-summary",
|
| 25 |
+
"Model_2": "pszemraj/long-t5-tglobal-base-sci-simplify",
|
| 26 |
+
"Model_3": "pszemraj/long-t5-tglobal-base-sci-simplify-elife",
|
| 27 |
+
"Model_4": "pszemraj/long-t5-tglobal-base-16384-booksci-summary-v1",
|
| 28 |
+
"Model_5": "pszemraj/pegasus-x-large-book-summary",
|
| 29 |
}
|
| 30 |
selected_model_identifier = MODEL_OPTIONS.get(model_name)
|
| 31 |
device = "cuda" if torch.cuda.is_available() else "cpu"
|