added generate.py

generate.py

@@ -0,0 +1,45 @@
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+tokenizer = AutoTokenizer.from_pretrained("tuanle/VN-News-GPT2", cache_dir="cache/")
+model = AutoModelForCausalLM.from_pretrained("tuanle/VN-News-GPT2", cache_dir="cache/").to(device)
+print("Loading model...")
+print("Model is ready to serve...")
+
+def generate(category, headline, 
+             min_len = 60, 
+             max_len = 768, 
+             num_beams = 5, 
+             num_return_sequences = 3,
+             top_k = 50,
+             top_p = 1):
+    """
+        top_p: If set to float < 1, only the most probable tokens with probabilities that add up to top_p or higher are kept for generation.
+        top_k: The number of highest probability vocabulary tokens to keep for top-k-filtering.
+        num_beams: Number of beams for beam search. 1 means no beam search.
+    """
+    text = f"<|startoftext|> {category} <|headline|> {headline}"
+
+    input_ids = tokenizer.encode(text, return_tensors='pt').to(device)
+
+    sample_outputs = model.generate(input_ids,
+                                    do_sample=True,
+                                    max_length=max_len,
+                                    min_length=min_len,
+                                    #    temperature = .8,
+                                    top_k= top_k,
+                                    top_p = top_p,
+                                    num_beams= num_beams,
+                                    early_stopping= True,
+                                    no_repeat_ngram_size= 2  ,
+                                    num_return_sequences= num_return_sequences)
+
+    outputs = []
+    for i, sample_output in enumerate(sample_outputs):
+        temp = tokenizer.decode(sample_output.tolist())
+        print(f">> Generated text {i+1}\n\n{temp}")
+        print('\n---')
+        outputs.append(temp)
+    return outputs