MeetJivani commited on
Commit
344edee
·
1 Parent(s): 36981a1

Update summarize.py

Browse files
Files changed (1) hide show
  1. 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
- "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"
 
21
  :return tuple: a tuple containing the model and tokenizer
22
  """
23
  MODEL_OPTIONS = {
24
+ "Text Summarization": "pszemraj/long-t5-tglobal-base-16384-book-summary",
25
+ "News Artical Summarization Alpha": "pszemraj/long-t5-tglobal-base-sci-simplify",
26
+ "News Artical Summarization Beta": "pszemraj/long-t5-tglobal-base-sci-simplify-elife",
27
+ "Scientific Document Summarization Alpha": "pszemraj/long-t5-tglobal-base-16384-booksci-summary-v1",
28
+ "Scientific Document Summarization Beta": "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"