Update app.py

app.py

@@ -10,7 +10,7 @@ b_tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom-1b1")#using small
 b_model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-1b1",device_map = "auto")
 
 g_tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b",token = hf_token)#using small paramerter version of model for faster inference on hf
-g_model = AutoModelForCausalLM.from_pretrained("google/gemma-2-2b",token = hf_token)
+g_model = AutoModelForCausalLM.from_pretrained("google/gemma-2-2b",token = hf_token,device_map="auto")
 
 def Sentence_Commpletion(model_name, input):
     
@@ -18,11 +18,11 @@ def Sentence_Commpletion(model_name, input):
     if model_name == "Bloom":
         tokenizer, model = b_tokenizer, b_model
         inputss = tokenizer(input, return_tensors="pt")
-        outputs = model.generate(inputss.input_ids, max_length=100, num_return_sequences=1)
+        outputs = model.generate(inputss.input_ids, max_new_tokens=31, num_return_sequences=1)
     elif model_name == "Gemma":
         tokenizer, model = g_tokenizer, g_model
-        input_ids = tokenizer(input, return_tensors="pt")
-        outputs = model.generate(**input_ids, max_new_tokens=32)
+        inputs= tokenizer(input, return_tensors="pt")
+        outputs = model.generate(inputs.input_ids, max_new_tokens=32)
     return tokenizer.decode(outputs[0])