Update app.py

app.py

@@ -1,100 +1,96 @@
-import json
-import requests
+import argparse
+import re
+import os
 import streamlit as st
 import random
-headers = {}
-MODELS = {
-    "GPT-2 Base": {
-        "url": "https://api-inference.huggingface.co/models/BigSalmon/InformalToFormalLincoln14"
-    }
-}
-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,
-            num_return_sequences:int,
-            top_p: float):
-    payload = {
-        "inputs": text,
-        "parameters": {
-            "max_new_tokens": max_len,
-            "top_k": top_k,
-            "top_p": top_p,
-            "temperature": temp,
-            "num_return_sequences": num_return_sequences,
-            "repetition_penalty": 2.0,
-        },
-        "options": {
-            "use_cache": True,
-        }
-    }
-    return query(payload, model_name)
-st.set_page_config(page_title="Thai GPT2 Demo")
-st.title("🐘 Thai GPT2")
-st.sidebar.subheader("Configurable parameters")
-max_len = st.sidebar.text_input(
-    "Maximum length",
-    value=5,
-    help="The maximum length of the sequence to be generated."
-)
-temp = st.sidebar.slider(
-    "Temperature",
-    value=1.0,
-    min_value=0.7,
-    max_value=1.5,
-    help="The value used to module the next token probabilities."
-)
-top_k = st.sidebar.text_input(
-    "Top k",
-    value=50,
-    help="The number of highest probability vocabulary tokens to keep for top-k-filtering."
-)
-num_return_sequences = st.sidebar.text_input(
-    "Returns",
-    value=5,
-    help="Number to return."
-)
-top_p = st.sidebar.text_input(
-    "Top p",
-    value=0.95,
-    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."
-)
-do_sample = st.sidebar.selectbox(
-    'Sampling?', (True, False), help="Whether or not to use sampling; use greedy decoding otherwise.")
-st.markdown(
-    """Thai GPT-2 demo. Part of the [Huggingface JAX/Flax event](https://discuss.huggingface.co/t/open-to-the-community-community-week-using-jax-flax-for-nlp-cv/)."""
-)
-model_name = st.selectbox('Model', (['GPT-2 Base']))
-hello = ['Custom', 'Yellow']
-prompt = st.selectbox('Prompt', options = hello)
-if prompt == "Custom":
-    prompt_box = "Enter your text here"
-text = st.text_area("Enter text", prompt_box)
-if st.button("Run"):
-    with st.spinner(text="Getting results..."):
-        st.subheader("Result")
-        print(f"maxlen:{max_len}, temp:{temp}, top_k:{top_k}, top_p:{top_p}")
-        result = process(text=text,
-                         model_name=model_name,
-                         max_len=int(max_len),
-                         temp=temp,
-                         num_return_sequences = int(num_return_sequences),
-                         top_k=int(top_k),
-                         top_p=float(top_p))           
-        st.write(result)
-        if "error" in result:
-            if type(result["error"]) is str:
-                st.write(f'{result["error"]}. 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:
-            print("hey")
+import numpy as np
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import tokenizers
+#os.environ["TOKENIZERS_PARALLELISM"] = "false"
+random.seed(None)
+suggested_text_list = ['פעם אחת, לפני שנים רבות','שלום, קוראים לי דורון ואני','בוקר טוב לכולם','ואז הפרתי את כל כללי הטקס כש']
+@st.cache(hash_funcs={tokenizers.Tokenizer: id, tokenizers.AddedToken: id})
+def load_model(model_name):
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    return model, tokenizer
+def extend(input_text, max_size=20, top_k=50, top_p=0.95):
+    if len(input_text) == 0:
+        input_text = ""
+    encoded_prompt = tokenizer.encode(
+    input_text, add_special_tokens=False, return_tensors="pt")
+    encoded_prompt = encoded_prompt.to(device)
+    if encoded_prompt.size()[-1] == 0:
+        input_ids = None
+    else:
+        input_ids = encoded_prompt
+    
+    output_sequences = model.generate(
+    input_ids=input_ids,
+    max_length=max_size + len(encoded_prompt[0]),
+    top_k=top_k, 
+    top_p=top_p, 
+    do_sample=True,
+    num_return_sequences=1)
+    # Remove the batch dimension when returning multiple sequences
+    if len(output_sequences.shape) > 2:
+        output_sequences.squeeze_()
+    generated_sequences = []
+    for generated_sequence_idx, generated_sequence in enumerate(output_sequences):        
+        generated_sequence = generated_sequence.tolist()
+        # Decode text
+        text = tokenizer.decode(generated_sequence, clean_up_tokenization_spaces=True)
+        # Remove all text after the stop token
+        text = text[: text.find(stop_token) if stop_token else None]
+        # Remove all text after 3 newlines
+        text = text[: text.find(new_lines) if new_lines else None]
+        # Add the prompt at the beginning of the sequence. Remove the excess text that was used for pre-processing
+        total_sequence = (
+            input_text + text[len(tokenizer.decode(encoded_prompt[0], clean_up_tokenization_spaces=True)) :]
+        )
+        generated_sequences.append(total_sequence)
+    
+    parsed_text = total_sequence.replace("<|startoftext|>", "").replace("\r","").replace("\n\n", "\n")
+    if len(parsed_text) == 0:
+        parsed_text = "שגיאה"
+    return parsed_text
+if __name__ == "__main__":
+    st.title("Hebrew GPT Neo (Small)")
+    pre_model_path = "Norod78/hebrew-gpt_neo-small"
+    model, tokenizer = load_model(pre_model_path)
+    stop_token = "<|endoftext|>"
+    new_lines = "\n\n\n"
+    np.random.seed(None)
+    random_seed = np.random.randint(10000,size=1)    
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    n_gpu = 0 if torch.cuda.is_available()==False else torch.cuda.device_count()
+    torch.manual_seed(random_seed)
+    if n_gpu > 0:
+        torch.cuda.manual_seed_all(random_seed)
+    model.to(device)
+    text_area = st.text_area("Enter the first few words (or leave blank), tap on \"Generate Text\" below. Tapping again will produce a different result.", 'האיש האחרון בעולם ישב לבד בחדרו כשלפתע נשמעה נקישה')
+    st.sidebar.subheader("Configurable parameters")
+    max_len = st.sidebar.slider("Max-Length", 0, 256, 192,help="The maximum length of the sequence to be generated.")
+    top_k = st.sidebar.slider("Top-K", 0, 100, 40, help="The number of highest probability vocabulary tokens to keep for top-k-filtering.")
+    top_p = st.sidebar.slider("Top-P", 0.0, 1.0, 0.92, 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.")
+    if st.button("Generate Text"):
+        with st.spinner(text="Generating results..."):
+            st.subheader("Result")
+            print(f"device:{device}, n_gpu:{n_gpu}, random_seed:{random_seed}, maxlen:{max_len}, top_k:{top_k}, top_p:{top_p}")
+            if len(text_area.strip()) == 0:
+                text_area = random.choice(suggested_text_list)
+            result = extend(input_text=text_area,                         
+                            max_size=int(max_len),                         
+                            top_k=int(top_k),
+                            top_p=float(top_p))
+            print("Done length: " + str(len(result)) + " bytes") 
+            #<div class="rtl" dir="rtl" style="text-align:right;">
+            st.markdown(f"<p dir=\"rtl\" style=\"text-align:right;\"> {result} </p>", unsafe_allow_html=True)
+            st.write("\n\nResult length: " + str(len(result)) + " bytes")
+            print(f"\"{result}\"")      
+    
+    st.markdown(    
+        """Hebrew text generation model (125M parameters) based on EleutherAI's gpt-neo architecture. Originally trained on a TPUv3-8 which was made avilable to me via the [TPU Research Cloud Program](https://sites.research.google/trc/)."""
+    )
+    st.markdown("<footer><hr><p style=\"font-size:14px\">Enjoy</p><p style=\"font-size:12px\">Created by <a href=\"https://linktr.ee/Norod78\">Doron Adler</a></p></footer> ", unsafe_allow_html=True)