Update inference.py

inference.py

@@ -4,18 +4,16 @@ import torch.nn.functional as F
 import pickle
 from safetensors.torch import load_file
 import logging
+import argparse
 
 # Set up logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
 # Hyperparameters
-embedding_dim = 256
-#change H number according to model used
+embedding_dim = 128
 hidden_dim = 256
-num_layers = 4
-sequence_length = 64
-temp = 1.0  # Temperature parameter
-top_k = 10  # Top-k sampling parameter
+num_layers = 2
+sequence_length = 10
 
 # LSTM Model
 class LSTMModel(nn.Module):
@@ -31,21 +29,6 @@ class LSTMModel(nn.Module):
         logits = self.fc(lstm_out[:, -1, :])
         return logits
 
-# Load the model and vocabulary
-logging.info('Loading the model and vocabulary...')
-model_state_dict = load_file('lstm_H256.safetensors')
-with open('word2idx.pkl', 'rb') as f:
-    word2idx = pickle.load(f)
-with open('idx2word.pkl', 'rb') as f:
-    idx2word = pickle.load(f)
-
-vocab_size = len(word2idx)
-model = LSTMModel(vocab_size, embedding_dim, hidden_dim, num_layers)
-model.load_state_dict(model_state_dict)
-model.eval()
-
-logging.info('Model and vocabulary loaded successfully.')
-
 # Function to predict the next word with temperature and top-k sampling
 def predict_next_word(model, word2idx, idx2word, sequence, sequence_length, temp, top_k):
     model.eval()
@@ -62,22 +45,52 @@ def predict_next_word(model, word2idx, idx2word, sequence, sequence_length, temp
         return predicted_word
 
 # Function to generate a sentence
-def generate_sentence(model, word2idx, idx2word, start_sequence, sequence_length, temp, top_k, max_length=50):
+def generate_sentence(model, word2idx, idx2word, start_sequence, sequence_length, temp, top_k, max_length):
     sentence = start_sequence
     for _ in range(max_length):
         next_word = predict_next_word(model, word2idx, idx2word, sentence, sequence_length, temp, top_k)
         sentence += ' ' + next_word
-        if next_word == '<pad>' or next_word == 'User':
+        if next_word == '<pad>':
             break
     return sentence
 
-# Example usage
-start_sequence = "User : What is the capital of France ? Bot :"
+# Parse command-line arguments
+def parse_args():
+    parser = argparse.ArgumentParser(description='LSTM Next Word Prediction Chatbot')
+    parser.add_argument('--temp', type=float, default=1.0, help='Temperature parameter')
+    parser.add_argument('--top_k', type=int, default=10, help='Top-k sampling parameter')
+    parser.add_argument('--model_file', type=str, default='lstm_model.safetensors', help='Path to the safetensors model file')
+    parser.add_argument('--start_sequence', type=str, default='Once upon a time', help='Starting sequence for sentence generation')
+    parser.add_argument('--max_length', type=int, default=50, help='Maximum number of words to generate')
+    return parser.parse_args()
+
+# Main function
+def main():
+    args = parse_args()
+    temp = args.temp
+    top_k = args.top_k
+    model_file = args.model_file
+    start_sequence = args.start_sequence
+    max_length = args.max_length
+
+    logging.info(f'Loading the model and vocabulary from {model_file}...')
+    model_state_dict = load_file(model_file)
+    with open('word2idx.pkl', 'rb') as f:
+        word2idx = pickle.load(f)
+
+    # Generate idx2word from word2idx
+    idx2word = {idx: word for word, idx in word2idx.items()}
+
+    vocab_size = len(word2idx)
+    model = LSTMModel(vocab_size, embedding_dim, hidden_dim, num_layers)
+    model.load_state_dict(model_state_dict)
+    model.eval()
 
+    logging.info('Model and vocabulary loaded successfully.')
+    logging.info(f'Starting sequence: {start_sequence}')
+    logging.info(f'Temperature: {temp}, Top-k: {top_k}, Max Length: {max_length}')
+    generated_sentence = generate_sentence(model, word2idx, idx2word, start_sequence, sequence_length, temp, top_k, max_length)
+    logging.info(f'Generated sentence: {generated_sentence}')
 
-temp = 0.5 # Adjust temperature
-top_k = 32 # Adjust top-k
-logging.info(f'Starting sequence: {start_sequence}')
-logging.info(f'Temperature: {temp}, Top-k: {top_k}')
-generated_sentence = generate_sentence(model, word2idx, idx2word, start_sequence, sequence_length, temp, top_k)
-logging.info(f'Generated sentence: {generated_sentence}')
+if __name__ == '__main__':
+    main()