Mohamad-Jaallouk commited on
Commit
47f34c8
1 Parent(s): 45aefee

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -6,14 +6,14 @@ import torch
6
 
7
 
8
  class ModelProcessor:
9
- device = "cuda:0" if torch.cuda.is_available() else "cpu"
10
  def __init__(self, repo_id="HuggingFaceTB/cosmo-1b"):
 
11
  # Initialize the tokenizer
12
  self.tokenizer = AutoTokenizer.from_pretrained(repo_id, use_fast=True)
13
 
14
  # Initialize and configure the model
15
  self.model = AutoModelForCausalLM.from_pretrained(
16
- repo_id, torch_dtype=torch.float16, device_map={"": device}, trust_remote_code=True
17
  )
18
  self.model.eval() # Set the model to evaluation mode
19
 
@@ -67,7 +67,7 @@ class ModelProcessor:
67
 
68
  processor = ModelProcessor()
69
 
70
- @spaces.GPU
71
  def detect(prompt):
72
  prediction = processor.process_data_and_compute_statistics(prompt)
73
  if prediction == 1:
 
6
 
7
 
8
  class ModelProcessor:
 
9
  def __init__(self, repo_id="HuggingFaceTB/cosmo-1b"):
10
+ self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
11
  # Initialize the tokenizer
12
  self.tokenizer = AutoTokenizer.from_pretrained(repo_id, use_fast=True)
13
 
14
  # Initialize and configure the model
15
  self.model = AutoModelForCausalLM.from_pretrained(
16
+ repo_id, torch_dtype=torch.float16, device_map={"": self.device}, trust_remote_code=True
17
  )
18
  self.model.eval() # Set the model to evaluation mode
19
 
 
67
 
68
  processor = ModelProcessor()
69
 
70
+ @spaces.GPU(duration=180)
71
  def detect(prompt):
72
  prediction = processor.process_data_and_compute_statistics(prompt)
73
  if prediction == 1: