Update app.py

app.py

@@ -1,16 +1,16 @@
-import os
+#import os
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
-hf_token = os.getenv("hftoken")
+#hf_token = os.getenv("hftoken")
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 
 
 
 b_tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom-560m")#using small parameter version of model for faster inference on hf
 b_model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-560m")
 
-g_tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b",use_auth_token = hf_token)#using small paramerter version of model for faster inference on hf
-g_model = AutoModelForCausalLM.from_pretrained("google/gemma-2-2b",use_auth_token = hf_token)
+g_tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b")#using small paramerter version of model for faster inference on hf
+g_model = AutoModelForCausalLM.from_pretrained("google/gemma-2-2b")
 
 def Sentence_Commpletion(model_name, input):
     
@@ -21,8 +21,8 @@ def Sentence_Commpletion(model_name, input):
         outputs = model.generate(inputss.input_ids, max_length=30, num_return_sequences=1)
     elif model_name == "Gemma":
         tokenizer, model = g_tokenizer, g_model
-        inputss = tokenizer(input, return_tensors="pt").to("cuda")
-        outputs = model.generate(inputss.input_ids, max_new_tokens=20).to("cuda")
+        input_ids = tokenizer(input, return_tensors="pt").to("cuda")
+        outputs = model.generate(**input_ids, max_new_tokens=32).to("cuda")
     return tokenizer.decode(outputs[0])