revert

app.py

@@ -4,33 +4,25 @@ from mtranslate import translate
 from prompts import PROMPT_LIST
 import streamlit as st
 import random
-import transformers
-from transformers import GPT2Tokenizer, GPT2LMHeadModel
 import fasttext
 import SessionState
 
+headers = {}
+
 LOGO = "huggingwayang.png"
 
 MODELS = {
-    "GPT-2 Small": "flax-community/gpt2-small-indonesian",
-    "GPT-2 Medium": "flax-community/gpt2-medium-indonesian",
-    "GPT-2 Small finetuned on Indonesian academic journals": "Galuh/id-journal-gpt2"
+    "GPT-2 Small": {
+        "url": "https://api-inference.huggingface.co/models/flax-community/gpt2-small-indonesian"
+    },
+    "GPT-2 Medium": {
+        "url": "https://api-inference.huggingface.co/models/flax-community/gpt2-medium-indonesian"
+    },
+    "GPT-2 Small finetuned on Indonesian academic journals": {
+        "url": "https://api-inference.huggingface.co/models/Galuh/id-journal-gpt2"
+    },
 }
 
-headers = {}
-
-@st.cache(show_spinner=False)
-def load_gpt(model_type):
-    model = GPT2LMHeadModel.from_pretrained(MODELS[model_type])
-
-    return model
-
-@st.cache(show_spinner=False, hash_funcs={transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer: lambda _: None})
-def load_gpt_tokenizer(model_type):
-    tokenizer = GPT2Tokenizer.from_pretrained(MODELS[model_type])
-
-    return tokenizer
-
 def get_image(text: str):
     url = "https://wikisearch.uncool.ai/get_image/"
     try:
@@ -46,10 +38,44 @@ def get_image(text: str):
         image = ""
     return image
 
+def query(payload, model_name):
+    data = json.dumps(payload)
+    # print("model url:", MODELS[model_name]["url"])
+    response = requests.request("POST", MODELS[model_name]["url"], headers=headers, data=data)
+    return json.loads(response.content.decode("utf-8"))
+
+def process(text: str,
+            model_name: str,
+            max_len: int,
+            temp: float,
+            top_k: int,
+            top_p: float):
+
+    payload = {
+        "inputs": text,
+        "parameters": {
+            "max_new_tokens": max_len,
+            "top_k": top_k,
+            "top_p": top_p,
+            "temperature": temp,
+            "repetition_penalty": 2.0,
+        },
+        "options": {
+            "use_cache": True,
+        }
+    }
+    return query(payload, model_name)
+
 st.set_page_config(page_title="Indonesian GPT-2 Demo")
 
 st.title("Indonesian GPT-2")
 
+try:
+    token = st.secrets["flax_community_token"]
+    headers = {"Authorization": f"Bearer {token}"}
+except FileNotFoundError:
+    print(f"Token is not found")
+
 ft_model = fasttext.load_model('lid.176.ftz')
 
 # Sidebar
@@ -138,41 +164,42 @@ if st.button("Run"):
             text = translate(session_state.text, "id", lang)
 
         st.subheader("Result")
-        model = load_gpt(model_name)
-        tokenizer = load_gpt_tokenizer(model_name)
-
-        input_ids = tokenizer.encode(text, return_tensors='pt')
-        output = model.generate(input_ids=input_ids,
-                                max_length=max_len,
-                                temperature=temp,
-                                top_k=top_k,
-                                top_p=top_p,
-                                repetition_penalty=2.0)
-
-        text = tokenizer.decode(output[0], 
-                                skip_special_tokens=True)
-        st.write(text.replace("\n", "  \n"))
-
-        st.text("Translation")
-        translation = translate(text, "en", "id")
-
-        if lang == "id":
-            st.write(translation.replace("\n", "  \n"))
-
-        else:
-            st.write(translate(text, lang, "id").replace("\n", "  \n"))
-
-        image_cat = "https://media.giphy.com/media/vFKqnCdLPNOKc/giphy.gif"
-        image = get_image(translation.replace("\"", "'"))
-
-        if image is not "":
-            st.image(image, width=400)
-
+        result = process(text=text,
+                         model_name=model_name,
+                         max_len=int(max_len),
+                         temp=temp,
+                         top_k=int(top_k),
+                         top_p=float(top_p))
+
+        # print("result:", result)
+        if "error" in result:
+            if type(result["error"]) is str:
+                st.write(f'{result["error"]}.', end=" ")
+                if "estimated_time" in result:
+                    st.write(f'Please try it again in about {result["estimated_time"]:.0f} seconds')
+            else:
+                if type(result["error"]) is list:
+                    for error in result["error"]:
+                        st.write(f'{error}')
         else:
-            # display cat image if no image found
-            st.image(image_cat, width=400)
+            result = result[0]["generated_text"]
+            st.write(result.replace("\n", "  \n"))
+            st.text("Translation")
+            translation = translate(result, "en", "id")
+            if lang == "id":
+                st.write(translation.replace("\n", "  \n"))
+            else:
+                st.write(translate(result, lang, "id").replace("\n", "  \n"))
+
+            image_cat = "https://media.giphy.com/media/vFKqnCdLPNOKc/giphy.gif"
+            image = get_image(translation.replace("\"", "'"))
+            if image is not "":
+                st.image(image, width=400)
+            else:
+                # display cat image if no image found
+                st.image(image_cat, width=400)
 
         # Reset state
         session_state.prompt = None
         session_state.prompt_box = None
-        session_state.text = None
+        session_state.text = None

requirements.txt

@@ -1,9 +1,5 @@
-transformers
 streamlit 
 requests==2.24.0
 requests-toolbelt==0.9.1
 mtranslate
--f https://download.pytorch.org/whl/torch_stable.html
-torch==1.7.1+cpu; sys_platform == 'linux'
-torch==1.7.1; sys_platform == 'darwin'
 fasttext