modified: app.py

new file: requirements.txt

app.py

@@ -1,11 +1,12 @@
-# git clone https://huggingface.co/Pipatpong/vcm_santa
+# git clone https://huggingface.co/spaces/Pipatpong/VCM_Demo
 
 import gradio as gr
 import re
+import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 checkpoint = "Pipatpong/vcm_santa"
-device = "cpu"
+device = "cuda" if torch.cuda.is_available() else "cpu"
 tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(checkpoint, trust_remote_code=True, device_map="auto", load_in_8bit=True)
 
@@ -13,7 +14,6 @@ def generate(text, max_length, num_return_sequences=1):
     inputs = tokenizer.encode(text, padding=False, add_special_tokens=False, return_tensors="pt")
     outputs = model.generate(inputs, max_length=max_length, num_return_sequences=num_return_sequences)
     gen_text = "Assignment : " + tokenizer.decode(outputs[0]).split("#")[0] if "#" else "Assignment : " + tokenizer.decode(outputs[0])
-
     return gen_text