Update app.py

app.py

@@ -1,4 +1,4 @@
-from transformers import RobertaTokenizer, RobertaForMaskedLM, GPT2Tokenizer, GPTJForCausalLM
+from transformers import RobertaTokenizer, RobertaForMaskedLM, GPT2Tokenizer, GPTNeoForCausalLM
 import torch
 import wikipedia
 import re
@@ -8,16 +8,12 @@ import syllables
 from aitextgen import aitextgen
 nltk.download('cmudict')
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
 masked_tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
 masked_model = RobertaForMaskedLM.from_pretrained('roberta-base')
 
 causal_tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
-gptj_tokenizer = GPT2Tokenizer.from_pretrained("EleutherAI/gpt-j-6B")
-gptj_model = GPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B", revision="float16", low_cpu_mem_usage=True)
-gptj_model.to(device)
-#gpt_neo = GPT2LMHeadModel.from_pretrained("gpt2", pad_token_id=causal_tokenizer.eos_token_id)
+gptneo_tokenizer = GPT2Tokenizer.from_pretrained("EleutherAI/gpt-neo-1.3B")
+gptneo_model = GPTNeoForCausalLM.from_pretrained("EleutherAI/gpt-neo-1.3B")
 
 # Without any parameters, aitextgen() will download, cache, and load the 124M GPT-2 "small" model
 gpt2 = aitextgen()
@@ -132,17 +128,17 @@ def get_rhyming_line(prompt, rhyming_word, inputs_len):
     print(f"Final Sentence: {final_sentence}")
     return final_sentence
 
-def gptj_summary(topic):
-    input_ids = gptj_tokenizer(f"Here is some information about {topic}", return_tensors="pt").input_ids.to(device)
-    generated_ids = gptj_model.generate(input_ids, do_sample=True, temperature=0.9, max_length=200)
-    generated_text = gptj_tokenizer.decode(generated_ids[0])
+def gptneo_summary(topic):
+    input_ids = gptneo_tokenizer(f"Here is some information about {topic}", return_tensors="pt").input_ids.to(device)
+    generated_ids = gptneo_model.generate(input_ids, do_sample=True, temperature=0.9, max_length=200)
+    generated_text = gptneo_tokenizer.decode(generated_ids[0])
     return generated_text 
     
 def generate(topic, wiki=True):
     if wiki:
         topic_summary = remove_punctuation(wikipedia.summary(topic))
     else:
-        topic_summary = remove_punctuation(gptj_summary(topic))
+        topic_summary = remove_punctuation(gptneo_summary(topic))
     word_list = topic_summary.split()
     topic_summary_len = len(topic_summary)
     no_of_words = len(word_list)
@@ -215,12 +211,12 @@ def generate(topic, wiki=True):
 
 def compare_summaries(topic):
     wiki_limerick = generate(topic, wiki=True)
-    gptj_limerick = generate(topic, wiki=False)
+    gptneo_limerick = generate(topic, wiki=False)
 
     output = f"Limerick with Wikipedia summary of topic as prompt: \n"
     output += wiki_limerick + "\n"
-    output += f"Limerick with GPT-J summary of topic as prompt: \n"
-    output += gptj_limerick
+    output += f"Limerick with GPT Neo summary of topic as prompt: \n"
+    output += gptneo_limerick
 
     return output