|
|
import os |
|
|
import torch |
|
|
import torch.nn.functional as F |
|
|
import tiktoken |
|
|
from model import ChatGCLM, MAX_SEQ_LEN |
|
|
|
|
|
MODEL_PATH = None |
|
|
for f in os.listdir("."): |
|
|
if f.startswith("ChatGCLM_") and f.endswith(".pt"): |
|
|
MODEL_PATH = f |
|
|
break |
|
|
|
|
|
if MODEL_PATH is None: |
|
|
print("Error: No model checkpoint found!") |
|
|
print("Please train the model first with: python3 train_chatgclm.py") |
|
|
exit(1) |
|
|
|
|
|
TOKENIZER_NAME = "gpt2" |
|
|
EOS_ID = 2 |
|
|
|
|
|
def load_model(device): |
|
|
tok = tiktoken.get_encoding(TOKENIZER_NAME) |
|
|
vocab_size = tok.n_vocab |
|
|
|
|
|
model = ChatGCLM(vocab_size).to(device) |
|
|
if os.path.exists(MODEL_PATH) and os.path.getsize(MODEL_PATH) > 0: |
|
|
print(f"Loading model from: {MODEL_PATH}") |
|
|
model.load_state_dict(torch.load(MODEL_PATH, map_location=device)) |
|
|
model.eval() |
|
|
return model, tok |
|
|
else: |
|
|
print(f"Error: Could not load model from {MODEL_PATH}") |
|
|
return None, None |
|
|
|
|
|
@torch.no_grad() |
|
|
def generate(model, prompt, tokenizer, device, max_new_tokens=200, temperature=0.8, top_k=50): |
|
|
model.eval() |
|
|
input_ids = tokenizer.encode(prompt) |
|
|
x = torch.tensor([input_ids], dtype=torch.long, device=device) |
|
|
|
|
|
print(f"\n{'='*70}") |
|
|
print(f"PROMPT: {prompt}") |
|
|
print(f"{'='*70}") |
|
|
print("GENERATED TEXT:") |
|
|
print(prompt, end="", flush=True) |
|
|
|
|
|
generated_tokens = [] |
|
|
for _ in range(max_new_tokens): |
|
|
ctx = x[:, -MAX_SEQ_LEN:] if x.size(1) > MAX_SEQ_LEN else x |
|
|
logits = model(ctx) |
|
|
next_token_logits = logits[:, -1, :] / temperature |
|
|
|
|
|
if top_k is not None: |
|
|
v, _ = torch.topk(next_token_logits, min(top_k, next_token_logits.size(-1))) |
|
|
next_token_logits[next_token_logits < v[:, [-1]]] = -float('Inf') |
|
|
|
|
|
probs = F.softmax(next_token_logits, dim=-1) |
|
|
next_token = torch.multinomial(probs, num_samples=1) |
|
|
idx = next_token.item() |
|
|
|
|
|
if idx == EOS_ID: |
|
|
break |
|
|
|
|
|
x = torch.cat((x, next_token), dim=1) |
|
|
generated_tokens.append(idx) |
|
|
token_text = tokenizer.decode([idx]) |
|
|
print(token_text, end="", flush=True) |
|
|
|
|
|
print(f"\n{'='*70}\n") |
|
|
return tokenizer.decode(generated_tokens) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" |
|
|
print(f"Using device: {device}") |
|
|
|
|
|
model, tokenizer = load_model(device) |
|
|
|
|
|
if model is None: |
|
|
exit(1) |
|
|
|
|
|
test_prompts = [ |
|
|
"Once upon a time", |
|
|
"The future of AI is", |
|
|
"In a world where", |
|
|
] |
|
|
|
|
|
print("\n" + "="*70) |
|
|
print("ChatGCLM Text Generation Demo") |
|
|
print("="*70) |
|
|
|
|
|
for prompt in test_prompts: |
|
|
generate(model, prompt, tokenizer, device, max_new_tokens=150, temperature=0.8, top_k=50) |
|
|
|
|
|
print("\n" + "="*70) |
|
|
print("Interactive Mode - Enter your own prompts!") |
|
|
print("="*70) |
|
|
|
|
|
while True: |
|
|
user_prompt = input("\nEnter prompt (or 'exit' to quit): ") |
|
|
if user_prompt.lower() == 'exit': |
|
|
break |
|
|
if user_prompt.strip(): |
|
|
generate(model, user_prompt, tokenizer, device, max_new_tokens=200, temperature=0.8, top_k=50) |
|
|
|
