Update app.py

app.py

@@ -26,17 +26,19 @@ import gradio as gr
 from transformers import GPTJForCausalLM
 import torch    
 
-model = GPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B", revision="float16", torch_dtype=torch.float16, low_cpu_mem_usage=False)
-tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B")
+tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom-560m")
+model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-560m")
 
 def get_result_with_bloom(text):
-  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
-
+  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