use news-api

app/app.py

@@ -4,11 +4,11 @@ from mtranslate import translate
 from prompts import PROMPT_LIST
 import random
 import time
-from transformers import pipeline, set_seed, AutoConfig, AutoTokenizer, GPT2LMHeadModel, GPT2Tokenizer
 import psutil
 import torch
 import os
-from abstract_dataset import AbstractDataset
+import requests
+
 
 # st.set_page_config(page_title="Indonesian GPT-2")
 
@@ -16,6 +16,7 @@ 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)
+news_api_auth_token = os.getenv("NEWS_API_AUTH_TOKEN", False)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 MODELS = {
@@ -59,51 +60,39 @@ ___
 model_type = st.sidebar.selectbox('Model', (MODELS.keys()))
 
 
-@st.cache(suppress_st_warning=True, allow_output_mutation=True)
-def get_generator(model_name: str):
-    st.write(f"Loading the GPT2 model {model_name}, please wait...")
-    tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=hf_auth_token)
-    model = GPT2LMHeadModel.from_pretrained(model_name, pad_token_id=tokenizer.eos_token_id, use_auth_token=hf_auth_token)
-    model.to(device)
-    model.resize_token_embeddings(len(tokenizer))
-    return model, tokenizer
-
-
 # Disable the st.cache for this function due to issue on newer version of streamlit
 # @st.cache(suppress_st_warning=True, hash_funcs={tokenizers.Tokenizer: id})
-def process(text_generator, tokenizer, title: str, keywords: str, text: str,
+def process(title: str, keywords: str, text: str,
             max_length: int = 250, do_sample: bool = True, top_k: int = 50, top_p: float = 0.95,
             temperature: float = 1.0, max_time: float = 120.0, seed=42, repetition_penalty=1.0,
             penalty_alpha = 0.6):
     # st.write("Cache miss: process")
-    set_seed(seed)
-    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)
-    prompt = f"title: {title}\nkeywords: {keywords}\n{text}"
-
-    generated = torch.tensor(tokenizer.encode(prompt)).unsqueeze(0)
-    generated = generated.to(device)
-
-    text_generator.eval()
-    sample_outputs = text_generator.generate(generated,
-                                             penalty_alpha=penalty_alpha,
-                                             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
-                                             )
-    result = tokenizer.decode(sample_outputs[0], skip_special_tokens=True)
-    prefix_length = len(title) + len(keywords) + len("title: keywords: ") + 2
-    result = result[prefix_length:]
-    title_index = result.find("title: ")
-    result = result[:title_index] if title_index > 0 else result
-    return result
+    url = 'https://news-api.uncool.ai/api/text_generator/v1'
+    print("news_api_auth_token:", news_api_auth_token)
+    headers = {'Authorization': 'Bearer ' + news_api_auth_token}
+    print("Requesting to news-api.uncool.ai with headers: ", headers)
+    data = {
+        "title": title,
+        "keywords": keywords,
+        "text": text,
+        "max_length": max_length,
+        "do_sample": do_sample,
+        "top_k": top_k,
+        "top_p": top_p,
+        "temperature": temperature,
+        "max_time": max_time,
+        "seed": seed,
+        "repetition_penalty": repetition_penalty,
+        "penalty_alpha": penalty_alpha
+    }
+    r = requests.post(url, headers=headers, data=data)
+    if r.status_code == 200:
+        result = r.json()['generated_text']
+        title_index = result.find("title: ")
+        result = result[:title_index] if title_index > 0 else result
+        return result
+    else:
+        return "Error: " + r.text
 
 
 st.title("Indonesian GPT-2 Applications")
@@ -215,10 +204,6 @@ if prompt_group_name in ["Indonesian Newspaper"]:
                 max_value=2.0
             )
 
-    for group_name in MODELS:
-        if MODELS[group_name]["group"] in ["Indonesian Newspaper"]:
-            MODELS[group_name]["text_generator"], MODELS[group_name]["tokenizer"] = \
-                get_generator(MODELS[group_name]["name"])
     # st.write(f"Generator: {MODELS}'")
     if st.button("Run"):
         with st.spinner(text="Getting results..."):
@@ -226,8 +211,7 @@ if prompt_group_name in ["Indonesian Newspaper"]:
             st.subheader("Result")
             time_start = time.time()
             # text_generator = MODELS[model_type]["text_generator"]
-            result = process(MODELS[model_type]["text_generator"], MODELS[model_type]["tokenizer"],
-                             title=session_state.title,
+            result = process(title=session_state.title,
                              keywords=session_state.keywords,
                              text=session_state.text, max_length=int(max_length),
                              temperature=temperature, do_sample=do_sample, penalty_alpha=penalty_alpha,