run

app/app.py

@@ -88,15 +88,10 @@ def process(text_generator, tokenizer, title: str, keywords: str, text: str,
         min_penalty = 1.05
         max_penalty = 1.5
         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)
-
-    special_tokens = AbstractDataset.special_tokens
-    prompt = special_tokens['bos_token'] + title + \
-             special_tokens['sep_token'] + keywords + special_tokens['sep_token'] + text
-
-    print(f"title: {title}, keywords: {keywords}, text: {text}")
+    print("title:", title)
+    print("keywords:", keywords)
+    prompt = f"title: {title}\nkeywords: {keywords}\n{text}"
+    print("prompt: ", prompt)
 
     generated = torch.tensor(tokenizer.encode(prompt)).unsqueeze(0)
     # device = torch.device("cuda")
@@ -115,7 +110,7 @@ def process(text_generator, tokenizer, title: str, keywords: str, text: str,
                                              )
     result = tokenizer.decode(sample_outputs[0], skip_special_tokens=True)
     print(f"result: {result}")
-    prefix_length = len(title) + len(keywords)
+    prefix_length = len(title) + len(keywords) + len("title: keywords: ") + 2
     result = result[prefix_length:]
     return result
 
@@ -231,7 +226,7 @@ if prompt_group_name in ["Indonesian Newspaper"]:
         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}'")
+    # st.write(f"Generator: {MODELS}'")
     if st.button("Run"):
         with st.spinner(text="Getting results..."):
             memory = psutil.virtual_memory()