Update app.py

app.py

@@ -1,13 +1,13 @@
 import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
-#torch.set_default_tensor_type(torch.cuda.FloatTensor)
+torch.set_default_tensor_type(torch.cuda.FloatTensor)
 
-def generate(text):
+def generate(prompt):
     tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom-1b1")
     model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-1b1", use_cache=True)
-    input_ids = tokenizer(text, return_tensors="pt").to(0)
-    sample = model.generate(**input_ids, max_length=300, num_beams = 2, num_beam_groups = 2, top_k=1, temperature=0.9, repetition_penalty = 2.0, diversity_penalty = 0.9)
+    input_ids = tokenizer(prompt, return_tensors="pt").to(0)
+    sample = model.generate(**input_ids, max_length=100, num_beams = 2, num_beam_groups = 2, top_k=1, temperature=0.9, repetition_penalty = 2.0, diversity_penalty = 0.9)
     return tokenizer.decode(sample[0], truncate_before_pattern=[r"\n\n^#", "^'''", "\n\n\n"])
     
 with gr.Blocks() as demo: