use newspaper modmel

app/app.py

@@ -10,18 +10,18 @@ import torch
 import os
 from abstract_dataset import AbstractDataset
 
-
 # st.set_page_config(page_title="Indonesian GPT-2")
 
-mirror_url = "https://abstract-generator.ai-research.id/"
+mirror_url = "https://news-generator.ai-research.id/"
 if "MIRROR_URL" in os.environ:
     mirror_url = os.environ["MIRROR_URL"]
+hf_auth_token = os.getenv("HF_AUTH_TOKEN", False)
 
 MODELS = {
-    "Indonesian Academic Journal - Indonesian GPT-2 Medium": {
-        "group": "Indonesian Journal",
-        "name": "cahya/abstract-generator",
-        "description": "Abstract Generator using Indonesian GPT-2 Medium.",
+    "Indonesian Newspaper - Indonesian GPT-2 Medium": {
+        "group": "Indonesian Newspaper",
+        "name": "ai-research-id/gpt2-medium-newspaper",
+        "description": "Newspaper Generator using Indonesian GPT-2 Medium.",
         "text_generator": None,
         "tokenizer": None
     },
@@ -85,7 +85,7 @@ def process(text_generator, tokenizer, title: str, keywords: str, text: str,
     if repetition_penalty == 0.0:
         min_penalty = 1.05
         max_penalty = 1.5
-        repetition_penalty = max(min_penalty + (1.0-temperature) * (max_penalty-min_penalty), 0.8)
+        repetition_penalty = max(min_penalty + (1.0 - temperature) * (max_penalty - min_penalty), 0.8)
 
     keywords = [keyword.strip() for keyword in keywords.split(",")]
     keywords = AbstractDataset.join_keywords(keywords, randomize=False)
@@ -102,15 +102,16 @@ def process(text_generator, tokenizer, title: str, keywords: str, text: str,
 
     text_generator.eval()
     sample_outputs = text_generator.generate(generated,
-                                    do_sample=do_sample,
-                                    min_length=200,
-                                    max_length=max_length,
-                                    top_k=top_k,
-                                    top_p=top_p,
-                                    temperature=temperature,
-                                    repetition_penalty=repetition_penalty,
-                                    num_return_sequences=1
-                                    )
+                                             do_sample=do_sample,
+                                             min_length=200,
+                                             max_length=max_length,
+                                             top_k=top_k,
+                                             top_p=top_p,
+                                             temperature=temperature,
+                                             repetition_penalty=repetition_penalty,
+                                             num_return_sequences=1,
+                                             hf_auth_token=hf_auth_token
+                                             )
     result = tokenizer.decode(sample_outputs[0], skip_special_tokens=True)
     print(f"result: {result}")
     prefix_length = len(title) + len(keywords)
@@ -127,9 +128,9 @@ model_name = f"Model name: [{MODELS[model_type]['name']}](https://huggingface.co
 st.markdown(model_name)
 if prompt_group_name in ["Indonesian GPT-2", "Indonesian Literature", "Indonesian Journal"]:
     session_state = SessionState.get(prompt=None, prompt_box=None, text=None)
-    ALL_PROMPTS = list(PROMPT_LIST[prompt_group_name].keys())+["Custom"]
+    ALL_PROMPTS = list(PROMPT_LIST[prompt_group_name].keys()) + ["Custom"]
 
-    prompt = st.selectbox('Prompt', ALL_PROMPTS, index=len(ALL_PROMPTS)-1)
+    prompt = st.selectbox('Prompt', ALL_PROMPTS, index=len(ALL_PROMPTS) - 1)
 
     # Update prompt
     if session_state.prompt is None:
@@ -160,6 +161,12 @@ if prompt_group_name in ["Indonesian GPT-2", "Indonesian Literature", "Indonesia
         help="The maximum length of the sequence to be generated."
     )
 
+    decoding_methods = st.sidebar.radio(
+        "Set the decoding methods:",
+        key="decoding",
+        options=["Beam Search", "Sampling", "Contrastive Search"],
+    )
+
     temperature = st.sidebar.slider(
         "Temperature",
         value=0.4,
@@ -167,15 +174,14 @@ if prompt_group_name in ["Indonesian GPT-2", "Indonesian Literature", "Indonesia
         max_value=2.0
     )
 
-    do_sample = st.sidebar.checkbox(
-        "Use sampling",
-        value=True
-    )
-
     top_k = 30
     top_p = 0.95
+    repetition_penalty = 0.0
 
-    if do_sample:
+    if decoding_methods == "Beam Search":
+        do_sample = False
+    elif decoding_methods == "Sampling":
+        do_sample = True
         top_k = st.sidebar.number_input(
             "Top k",
             value=top_k,
@@ -187,6 +193,19 @@ if prompt_group_name in ["Indonesian GPT-2", "Indonesian Literature", "Indonesia
             help="If set to float < 1, only the most probable tokens with probabilities that add up to top_p or higher "
                  "are kept for generation."
         )
+    else:
+        do_sample = False
+        repetition_penalty = 1.0
+        penalty_alpha = st.sidebar.number_input(
+            "Penalty alpha",
+            value=0.6,
+            help="The penalty alpha for contrastive search."
+        )
+        top_k = st.sidebar.number_input(
+            "Top k",
+            value=4,
+            help="The number of highest probability vocabulary tokens to keep for top-k-filtering."
+        )
 
     seed = st.sidebar.number_input(
         "Random Seed",
@@ -194,22 +213,21 @@ if prompt_group_name in ["Indonesian GPT-2", "Indonesian Literature", "Indonesia
         help="The number used to initialize a pseudorandom number generator"
     )
 
-    repetition_penalty = 0.0
-    automatic_repetition_penalty = st.sidebar.checkbox(
-        "Automatic Repetition Penalty",
-        value=True
-    )
-
-    if not automatic_repetition_penalty:
-        repetition_penalty = st.sidebar.slider(
-            "Repetition Penalty",
-            value=1.0,
-            min_value=1.0,
-            max_value=2.0
+    if decoding_methods != "Contrastive Search":
+        automatic_repetition_penalty = st.sidebar.checkbox(
+            "Automatic Repetition Penalty",
+            value=True
         )
+        if not automatic_repetition_penalty:
+            repetition_penalty = st.sidebar.slider(
+                "Repetition Penalty",
+                value=1.0,
+                min_value=1.0,
+                max_value=2.0
+            )
 
     for group_name in MODELS:
-        if MODELS[group_name]["group"] in ["Indonesian GPT-2", "Indonesian Literature", "Indonesian Journal"]:
+        if MODELS[group_name]["group"] in ["Indonesian Newspaper"]:
             MODELS[group_name]["text_generator"], MODELS[group_name]["tokenizer"] = \
                 get_generator(MODELS[group_name]["name"])
 
@@ -226,15 +244,15 @@ if prompt_group_name in ["Indonesian GPT-2", "Indonesian Literature", "Indonesia
                              temperature=temperature, do_sample=do_sample,
                              top_k=int(top_k), top_p=float(top_p), seed=seed, repetition_penalty=repetition_penalty)
             time_end = time.time()
-            time_diff = time_end-time_start
-            #result = result[0]["generated_text"]
+            time_diff = time_end - time_start
+            # result = result[0]["generated_text"]
             st.write(result.replace("\n", "  \n"))
             st.text("Translation")
             translation = translate(result, "en", "id")
             st.write(translation.replace("\n", "  \n"))
             # st.write(f"*do_sample: {do_sample}, top_k: {top_k}, top_p: {top_p}, seed: {seed}*")
             info = f"""
-            *Memory: {memory.total/(1024*1024*1024):.2f}GB, used: {memory.percent}%, available: {memory.available/(1024*1024*1024):.2f}GB*        
+            *Memory: {memory.total / (1024 * 1024 * 1024):.2f}GB, used: {memory.percent}%, available: {memory.available / (1024 * 1024 * 1024):.2f}GB*        
             *Text generated in {time_diff:.5} seconds*
             """
             st.write(info)