Spaces:
Runtime error
Runtime error
Update app.py
Browse files
app.py
CHANGED
|
@@ -106,7 +106,7 @@ class TextGeneration:
|
|
| 106 |
set_seed(42)
|
| 107 |
|
| 108 |
def load(self):
|
| 109 |
-
logger.info("Loading model"
|
| 110 |
self.tokenizer = AutoTokenizer.from_pretrained(
|
| 111 |
self.model_name_or_path, use_auth_token=HF_AUTH_TOKEN if HF_AUTH_TOKEN else None,
|
| 112 |
)
|
|
@@ -118,7 +118,7 @@ class TextGeneration:
|
|
| 118 |
_ = self.model.eval()
|
| 119 |
device_number = -1 if DEVICE == "cpu" else int(DEVICE.split(":")[-1])
|
| 120 |
self.generator = pipeline(self.task, model=self.model, tokenizer=self.tokenizer, device=device_number)
|
| 121 |
-
logger.info("
|
| 122 |
# with torch.no_grad():
|
| 123 |
# tokens = tokenizer.encode(prompt, return_tensors='pt').to(device=device, non_blocking=True)
|
| 124 |
# gen_tokens = self.model.generate(tokens, do_sample=True, temperature=0.8, max_length=128)
|
|
|
|
| 106 |
set_seed(42)
|
| 107 |
|
| 108 |
def load(self):
|
| 109 |
+
logger.info("Loading model...")
|
| 110 |
self.tokenizer = AutoTokenizer.from_pretrained(
|
| 111 |
self.model_name_or_path, use_auth_token=HF_AUTH_TOKEN if HF_AUTH_TOKEN else None,
|
| 112 |
)
|
|
|
|
| 118 |
_ = self.model.eval()
|
| 119 |
device_number = -1 if DEVICE == "cpu" else int(DEVICE.split(":")[-1])
|
| 120 |
self.generator = pipeline(self.task, model=self.model, tokenizer=self.tokenizer, device=device_number)
|
| 121 |
+
logger.info("Loading model done.")
|
| 122 |
# with torch.no_grad():
|
| 123 |
# tokens = tokenizer.encode(prompt, return_tensors='pt').to(device=device, non_blocking=True)
|
| 124 |
# gen_tokens = self.model.generate(tokens, do_sample=True, temperature=0.8, max_length=128)
|