Update inference_fine_tune.py

inference_fine_tune.py

@@ -7,6 +7,15 @@ from pathlib import Path
 from config import get_config, get_weights_file_path
 from train import get_model
 
+def get_tokenizer(config)->Tokenizer:
+    tokenizers_path = Path(config['tokenizer_file'])
+    if Path.exists(tokenizers_path):
+        print("Loading tokenizer from ", tokenizers_path)
+        tokenizer = Tokenizer.from_file(str(tokenizers_path))
+        return tokenizer
+    else:
+        raise FileNotFoundError("Cant find tokenizer file : ",tokenizers_path)
+
 
 config = get_config("./openweb.config.json")
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -22,101 +31,6 @@ 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
-):
-    eos_idx = tokenizer.token_to_id('</s>')
-    pad_idx = tokenizer.token_to_id('<pad>')
-
-    # Start with the input text as initial decoder input
-    decoder_input = text.to(device)
-    if decoder_input.dim() == 1:
-       decoder_input = decoder_input.unsqueeze(0)
-
-
-    # Print the initial prompt
-
-    while decoder_input.shape[1] < 2000 :
-        # Apply causal mask based on current decoder_input length
-        # decoder_mask = (decoder_input != pad_idx).unsqueeze(0).int() & causal_mask(decoder_input.size(1)).type_as(mask).to(device)
-
-        # Get model output
-        out = model.decode(decoder_input)
-        logits = model.project(out[:, -1])  # Get logits for last token
-
-        # Sampling: temperature + top-k
-        logits = logits / temperature
-        top_k_logits, top_k_indices = torch.topk(logits, top_k)
-        probs = torch.softmax(top_k_logits, dim=-1)
-        next_token = torch.multinomial(probs, num_samples=1)
-        next_token = top_k_indices.gather(-1, next_token)
-
-        # Decode and print token
-        word = tokenizer.decode([next_token.item()])
-        print(word, end="", flush=True)
-
-        # Append next token
-
-        decoder_input = torch.cat([decoder_input, next_token], dim=1)
-        if decoder_input.shape[1] > max_len:
-            decoder_input = decoder_input[:,-max_len:]
-        
-        
-        if next_token.item() == eos_idx:
-            break
-
-    print()
-    return decoder_input.squeeze(0)
-
-
-
-def get_tokenizer(config)->Tokenizer:
-    tokenizers_path = Path(config['tokenizer_file'])
-    if Path.exists(tokenizers_path):
-        print("Loading tokenizer from ", tokenizers_path)
-        tokenizer = Tokenizer.from_file(str(tokenizers_path))
-        return tokenizer
-    else:
-        raise FileNotFoundError("Cant find tokenizer file : ",tokenizers_path)
-    
-def run_model(config):
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    print(f"Using device : {device}")
-    tokenizer = get_tokenizer(config)
-    model  = get_model(config, tokenizer.get_vocab_size()).to(device)
-    model_path = get_weights_file_path(config,config['preload'])
-    model.eval()
-    
-    if Path.exists(Path(model_path)):
-        print("Loading Model from : ", model_path)
-        state = torch.load(model_path)
-        model.load_state_dict(state['model_state_dict'])
-        print("You : ",end="")
-        input_text = input()
-        pad_token_id = tokenizer.token_to_id("<pad>")
-        user_token_id = tokenizer.token_to_id("<user>") 
-        ai_token_id = tokenizer.token_to_id("<ai>")
-        while input_text != "exit":
-            input_tokens = tokenizer.encode(input_text).ids[:-1]
-            input_tokens.extend([user_token_id] + input_tokens + [ai_token_id] )
-            
-            if len(input_tokens) > config['seq_len']:
-                print(f"exceeding max length of input : {config['seq_len']}")
-                continue
-            # if len(input_tokens) < config['seq_len']:
-                # input_tokens += [pad_token_id] * (config['seq_len'] - len(input_tokens))
-            input_tokens = torch.tensor(input_tokens)
-            output_tokens = generate_text(model, input_tokens, tokenizer, config['seq_len'], device )
-            print("MODEL : ",output_tokens)
-            output_text = tokenizer.decode(output_tokens.detach().cpu().numpy())
-            # print("Model : "+output_text)
-            print("You : ",end="")
-            input_text = input()
-            
-    else:
-        raise FileNotFoundError("Model File not found : "+ model_path)
-    
 def generate_response(prompt:str):
     print("Prompt : ",prompt)