Add inference script
Browse files- hindi_embeddings.py +3 -9
hindi_embeddings.py
CHANGED
|
@@ -384,24 +384,18 @@ class SentenceEmbeddingModel(nn.Module):
|
|
| 384 |
return pooled_output
|
| 385 |
|
| 386 |
class HindiEmbedder:
|
| 387 |
-
def __init__(self, model_path="/home/ubuntu/output/hindi-embeddings-custom-tokenizer/final"
|
| 388 |
"""
|
| 389 |
Initialize the Hindi sentence embedder.
|
| 390 |
|
| 391 |
Args:
|
| 392 |
model_path: Path to the model directory
|
| 393 |
-
tokenizer_path: Optional path to tokenizer. If None, will look in the model directory.
|
| 394 |
"""
|
| 395 |
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 396 |
print(f"Using device: {self.device}")
|
| 397 |
|
| 398 |
-
# Load tokenizer
|
| 399 |
-
|
| 400 |
-
# Try standard location in model directory
|
| 401 |
-
tokenizer_path = os.path.join(model_path, "tokenizer.model")
|
| 402 |
-
if not os.path.exists(tokenizer_path):
|
| 403 |
-
# Try original location
|
| 404 |
-
tokenizer_path = "/home/ubuntu/hindi_tokenizer/tokenizer.model"
|
| 405 |
|
| 406 |
if not os.path.exists(tokenizer_path):
|
| 407 |
raise FileNotFoundError(f"Could not find tokenizer at {tokenizer_path}")
|
|
|
|
| 384 |
return pooled_output
|
| 385 |
|
| 386 |
class HindiEmbedder:
|
| 387 |
+
def __init__(self, model_path="/home/ubuntu/output/hindi-embeddings-custom-tokenizer/final"):
|
| 388 |
"""
|
| 389 |
Initialize the Hindi sentence embedder.
|
| 390 |
|
| 391 |
Args:
|
| 392 |
model_path: Path to the model directory
|
|
|
|
| 393 |
"""
|
| 394 |
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 395 |
print(f"Using device: {self.device}")
|
| 396 |
|
| 397 |
+
# Load tokenizer - look for it in the model directory
|
| 398 |
+
tokenizer_path = os.path.join(model_path, "tokenizer.model")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 399 |
|
| 400 |
if not os.path.exists(tokenizer_path):
|
| 401 |
raise FileNotFoundError(f"Could not find tokenizer at {tokenizer_path}")
|