Update app.py

app.py

@@ -12,7 +12,54 @@ from peft import PeftModel
 
 from share_btn import community_icon_html, loading_icon_html, share_js, share_btn_css
 
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+FORMATS = """## Model Formats
+
+The model is pretrained on code and is formatted with special tokens in addition to the pure code data,\
+such as prefixes specifying the source of the file or tokens separating code from a commit message.\
+Use these templates to explore the model's capacities:
+
+### 1. Prefixes 🏷️
+For pure code files, use any combination of the following prefixes:
+
+```
+<reponame>REPONAME<filename>FILENAME<gh_stars>STARS\ncode<|endoftext|>
+```
+STARS can be one of: 0, 1-10, 10-100, 100-1000, 1000+
+
+### 2. Commits 💾
+The commits data is formatted as follows:
+
+```
+<commit_before>code<commit_msg>text<commit_after>code<|endoftext|>
+```
+
+### 3. Jupyter Notebooks 📓
+The model is trained on Jupyter notebooks as Python scripts and structured formats like:
+
+```
+<start_jupyter><jupyter_text>text<jupyter_code>code<jupyter_output>output<jupyter_text>
+```
+
+### 4. Issues 🐛
+We also trained on GitHub issues using the following formatting:
+```
+<issue_start><issue_comment>text<issue_comment>...<issue_closed>
+```
+
+### 5. Fill-in-the-middle 🧩
+Fill in the middle requires rearranging the model inputs. The playground handles this for you - all you need is to specify where to fill:
+```
+code before<FILL_HERE>code after
+```
+"""
+
+if not torch.cuda.is_available():
+    FORMATS += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
+
+if torch.cuda.is_available():
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    print(device)
 
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 CHECKPOINT_URL = "Salesforce/codegen-350M-mono"
@@ -86,53 +133,6 @@ model_map = {
 }
 
 
-FIM_PREFIX = "<fim_prefix>"
-FIM_MIDDLE = "<fim_middle>"
-FIM_SUFFIX = "<fim_suffix>"
-
-FIM_INDICATOR = "<FILL_HERE>"
-
-FORMATS = """## Model Formats
-
-The model is pretrained on code and is formatted with special tokens in addition to the pure code data,\
-such as prefixes specifying the source of the file or tokens separating code from a commit message.\
-Use these templates to explore the model's capacities:
-
-### 1. Prefixes 🏷️
-For pure code files, use any combination of the following prefixes:
-
-```
-<reponame>REPONAME<filename>FILENAME<gh_stars>STARS\ncode<|endoftext|>
-```
-STARS can be one of: 0, 1-10, 10-100, 100-1000, 1000+
-
-### 2. Commits 💾
-The commits data is formatted as follows:
-
-```
-<commit_before>code<commit_msg>text<commit_after>code<|endoftext|>
-```
-
-### 3. Jupyter Notebooks 📓
-The model is trained on Jupyter notebooks as Python scripts and structured formats like:
-
-```
-<start_jupyter><jupyter_text>text<jupyter_code>code<jupyter_output>output<jupyter_text>
-```
-
-### 4. Issues 🐛
-We also trained on GitHub issues using the following formatting:
-```
-<issue_start><issue_comment>text<issue_comment>...<issue_closed>
-```
-
-### 5. Fill-in-the-middle 🧩
-Fill in the middle requires rearranging the model inputs. The playground handles this for you - all you need is to specify where to fill:
-```
-code before<FILL_HERE>code after
-```
-"""
-
 theme = gr.themes.Monochrome(
     primary_hue="indigo",
     secondary_hue="blue",
@@ -182,9 +182,9 @@ def generate(
     else:
         output = ""
     # model.to(device)
-    input_ids = tokenizer(prompt, return_tensors="pt").to(device)
-    # generated_ids = model.generate(**input_ids, **generate_kwargs)
-    generated_ids = model.generate(**input_ids)
+    input_ids = tokenizer(prompt, return_tensors="pt").to(model.device)
+    # generated_ids = model.generate(**input_ids
+    generated_ids = model.generate(**input_ids, **generate_kwargs)
     
     return tokenizer.decode(generated_ids[0][input_ids["input_ids"].shape[1]:], skip_special_tokens=True).strip()