Update app.py

app.py

@@ -23,20 +23,19 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
 
-tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom-560m")
+from transformers import GPTJForCausalLM
+import torch    
 
-model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-560m")
+model = GPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B", revision="float16", torch_dtype=torch.float16, low_cpu_mem_usage=True)
+tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B")
 
 def get_result_with_bloom(text):
-  result_length = 200
-  inputs1 = tokenizer(text, return_tensors="pt")
-  output1 = tokenizer.decode(model.generate(inputs1["input_ids"],
-                       max_length=result_length, 
-                       num_beams=2, 
-                       no_repeat_ngram_size=2,
-                       early_stopping=True
-                      )[0])
-  return output1
+  context = text
+
+  input_ids = tokenizer(context, return_tensors="pt").input_ids
+  gen_tokens = model.generate(input_ids, do_sample=True, temperature=0.9, max_length=100,)
+  gen_text = tokenizer.batch_decode(gen_tokens)[0]
+  return gen_text