Spaces:
Runtime error
Runtime error
Mohamad-Jaallouk
commited on
Commit
•
47f34c8
1
Parent(s):
45aefee
Update app.py
Browse files
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:
|