Update app.py

app.py

@@ -4,13 +4,13 @@ from transformers import pipeline
 import os
 import torch
 
-description = """# <p style="text-align: center; color: white;"> 🎅 <span style='color: #ff75b3;'>SantaCoder:</span> Code Generation </p>
+description = """# <p style="text-align: center; color: white;"> 🎅 <span style='color: #ff75b3;'>SantaFixer:</span> Code Generation </p>
 <span style='color: white;'>This is a demo to generate code with <a href="https://huggingface.co/bigcode/santacoder" style="color: #ff75b3;">SantaCoder</a>,
 a 1.1B parameter model for code generation in Python, Java & JavaScript. The model can also do infilling, just specify where you would like the model to complete code
 with the <span style='color: #ff75b3;'>&lt;FILL-HERE&gt;</span> token.</span>"""
 
 token = os.environ["HUB_TOKEN"]
-device="cuda:0"
+device=torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
 
 
 FIM_PREFIX = "<fim-prefix>"
@@ -21,7 +21,7 @@ EOD = "<|endoftext|>"
 
 GENERATION_TITLE= "<p style='font-size: 16px; color: white;'>Generated code:</p>"
 
-tokenizer_fim = AutoTokenizer.from_pretrained("bigcode/santacoder", use_auth_token=token, padding_side="left")
+tokenizer_fim = AutoTokenizer.from_pretrained("lambdasec/santafixer", use_auth_token=token, padding_side="left")
 
 tokenizer_fim.add_special_tokens({
   "additional_special_tokens": [EOD, FIM_PREFIX, FIM_MIDDLE, FIM_SUFFIX, FIM_PAD],