Update app.py

app.py

@@ -7,11 +7,11 @@ import random
 # import torch
 
 # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model_name = "bloom-560m"
-model = AutoModelForCausalLM.from_pretrained(f'jslin09/{model_name}-finetuned-fraud').to('cuda')
-tokenizer = BloomTokenizerFast.from_pretrained(f'bigscience/{model_name}', bos_token = '<s>', eos_token = '</s>', pad_token = '<pad>')
-#model = AutoModelForCausalLM.from_pretrained("jslin09/gemma2-2b-fraud").to('cuda')
-#tokenizer = AutoTokenizer.from_pretrained("jslin09/gemma2-2b-fraud", bos_token = '<bos>', eos_token = '<eos>', pad_token = '<pad>')
+#model_name = "bloom-560m"
+#model = AutoModelForCausalLM.from_pretrained(f'jslin09/{model_name}-finetuned-fraud').to('cuda')
+#tokenizer = BloomTokenizerFast.from_pretrained(f'bigscience/{model_name}', bos_token = '<s>', eos_token = '</s>', pad_token = '<pad>')
+model = AutoModelForCausalLM.from_pretrained("jslin09/gemma2-2b-fraud").to('cuda')
+tokenizer = AutoTokenizer.from_pretrained("jslin09/gemma2-2b-fraud", bos_token = '<bos>', eos_token = '<eos>', pad_token = '<pad>')
 
 @spaces.GPU
 def rnd_generate(prompt):
@@ -45,6 +45,7 @@ def generate(prompt):
                        top_p=0.9
                       )
     try:
+        print(results)
         result = tokenizer.decode(results[0]).split("<bos>")[1].split("<eos>")[0]
     except IndexError:
         result = tokenizer.decode(results[0]).split("<bos>")[1]