Update inference_fine_tune.py

inference_fine_tune.py

@@ -7,6 +7,21 @@ from pathlib import Path
 from config import get_config, get_weights_file_path
 from train import get_model
 
+
+config = get_config("./openweb.config.json")
+device = "cuda" if torch.cuda.is_available() else "cpu"
+tokenizer = get_tokenizer(config)
+pad_token_id = tokenizer.token_to_id("<pad>")
+eos_token_id = tokenizer.token_to_id("</s>")
+user_token_id = tokenizer.token_to_id("<user>") 
+ai_token_id = tokenizer.token_to_id("<ai>")
+
+model  = get_model(config, tokenizer.get_vocab_size()).to(device)
+model_path = get_weights_file_path(config,config['preload'])
+model.eval()        
+state = torch.load(model_path,map_location=torch.device('cpu'))
+model.load_state_dict(state['model_state_dict'])
+
 def generate_text(
     model, text, tokenizer, max_len, device, 
     temperature=0.7, top_k=50
@@ -104,19 +119,7 @@ def run_model(config):
     
 def generate_response(prompt:str):
     print("Prompt : ",prompt)
-    config = get_config("./openweb.config.json")
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    tokenizer = get_tokenizer(config)
-    pad_token_id = tokenizer.token_to_id("<pad>")
-    eos_token_id = tokenizer.token_to_id("</s>")
-    user_token_id = tokenizer.token_to_id("<user>") 
-    ai_token_id = tokenizer.token_to_id("<ai>")
-
-    model  = get_model(config, tokenizer.get_vocab_size()).to(device)
-    model_path = get_weights_file_path(config,config['preload'])
-    model.eval()        
-    state = torch.load(model_path,map_location=torch.device('cpu'))
-    model.load_state_dict(state['model_state_dict'])
+    
     word = ""
     input_tokens = tokenizer.encode(prompt).ids
     input_tokens.extend([user_token_id] + input_tokens + [ai_token_id] )
@@ -149,8 +152,4 @@ def generate_response(prompt:str):
         if next_token.item() == eos_token_id:
             break
     print("Output : ",word)
-    return word
-
-if __name__ == "__main__":
-    config = get_config("openweb.config.json")
-    run_model(config)
+    return word