rootxhacker commited on
Commit
bcc9c6b
1 Parent(s): ccad9ef

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -33,7 +33,7 @@ tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
33
  # Load the Lora model
34
  model = PeftModel.from_pretrained(model, peft_model_id)
35
 
36
- @spaces.GPU(duration=200)
37
  def get_completion(query, model, tokenizer):
38
  try:
39
  # Move model to CUDA
@@ -56,7 +56,7 @@ def get_completion(query, model, tokenizer):
56
  model = model.cpu()
57
  torch.cuda.empty_cache()
58
 
59
- @spaces.GPU(duration=200)
60
  def code_review(code_to_analyze):
61
  query = f"As a code review expert, examine the following code for potential security flaws and provide guidance on secure coding practices:\n{code_to_analyze}"
62
  result = get_completion(query, model, tokenizer)
 
33
  # Load the Lora model
34
  model = PeftModel.from_pretrained(model, peft_model_id)
35
 
36
+ @spaces.GPU()
37
  def get_completion(query, model, tokenizer):
38
  try:
39
  # Move model to CUDA
 
56
  model = model.cpu()
57
  torch.cuda.empty_cache()
58
 
59
+ @spaces.GPU()
60
  def code_review(code_to_analyze):
61
  query = f"As a code review expert, examine the following code for potential security flaws and provide guidance on secure coding practices:\n{code_to_analyze}"
62
  result = get_completion(query, model, tokenizer)