File size: 6,601 Bytes
d757a4a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 | # use.py - Chat with VDrontMoE model (text continuation)
import torch
import argparse
import torch.nn.functional as F
import os
from transformers import AutoTokenizer
from create import VDrontMoEConfig, VDrontMoEModel
# --- Configuration ---
MODEL_DIR = "./VDrontMoE-2m-5e"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
TEMPERATURE = 0.4
MAX_NEW_TOKENS = 100
TOP_P = 0.9
TOP_K = 50
REPETITION_PENALTY = 1.1
# --- Load model ---
print(f"Loading model from {MODEL_DIR}...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
model = VDrontMoEModel.from_pretrained(MODEL_DIR).to(DEVICE)
model.eval()
print(f"Model loaded! Temperature: {TEMPERATURE}")
def generate(prompt: str, max_tokens: int = MAX_NEW_TOKENS) -> str:
"""Continue text from prompt."""
# Tokenize prompt
input_ids = tokenizer.encode(prompt, return_tensors="pt").to(DEVICE)
generated = input_ids.clone()
with torch.no_grad():
for _ in range(max_tokens):
# Truncate if too long
if generated.shape[1] > 512:
generated = generated[:, -512:]
# Forward pass
outputs = model(generated)
logits = outputs["logits"][:, -1, :] # Take last token
# Temperature
logits = logits / TEMPERATURE
# Top-K filter
if TOP_K > 0:
top_k_values, _ = torch.topk(logits, min(TOP_K, logits.size(-1)))
logits[logits < top_k_values[:, -1:]] = float('-inf')
# Top-P (nucleus) filter
if TOP_P < 1.0:
sorted_logits, sorted_indices = torch.sort(logits, descending=True)
cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
# Remove tokens with cumulative prob > TOP_P
sorted_indices_to_remove = cumulative_probs > TOP_P
sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()
sorted_indices_to_remove[:, 0] = False
indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
logits[indices_to_remove] = float('-inf')
# Repetition penalty
if REPETITION_PENALTY != 1.0:
for token_id in set(generated[0].tolist()[-10:]): # Last 10 tokens
if logits[0, token_id] > 0:
logits[0, token_id] /= REPETITION_PENALTY
else:
logits[0, token_id] *= REPETITION_PENALTY
# Sample next token
probs = F.softmax(logits, dim=-1)
next_token = torch.multinomial(probs, num_samples=1)
# Add to generated
generated = torch.cat([generated, next_token], dim=-1)
# Check for end of text
if next_token.item() == tokenizer.eos_token_id:
break
# Decode
full_text = tokenizer.decode(generated[0], skip_special_tokens=True)
return full_text
def chat():
"""Interactive text continuation mode."""
print("\n" + "=" * 60)
print("VDrontMoE-2m-5e - Text Continuation Mode")
print("=" * 60)
print("Commands:")
print(" /temp <value> - Set temperature (0.1 - 2.0)")
print(" /max <value> - Set max new tokens (10 - 500)")
print(" /clear - Clear screen")
print(" /exit - Exit")
print("=" * 60)
global TEMPERATURE, MAX_NEW_TOKENS
while True:
try:
# Get prompt
prompt = input("\nPrompt: ").strip()
# Commands
if prompt.startswith("/"):
parts = prompt.split()
cmd = parts[0]
if cmd == "/exit":
print("Goodbye!")
break
elif cmd == "/clear":
os.system('cls' if os.name == 'nt' else 'clear')
continue
elif cmd == "/temp" and len(parts) > 1:
TEMPERATURE = float(parts[1])
print(f"Temperature set to {TEMPERATURE}")
continue
elif cmd == "/max" and len(parts) > 1:
MAX_NEW_TOKENS = int(parts[1])
print(f"Max tokens set to {MAX_NEW_TOKENS}")
continue
else:
print("Unknown command")
continue
if not prompt:
print("Please enter a prompt!")
continue
# Generate
print("\nGenerating...")
result = generate(prompt)
# Display
print("\n" + "=" * 60)
print("GENERATED TEXT:")
print("=" * 60)
print(result)
print("=" * 60)
# Statistics
generated_part = result[len(prompt):]
new_tokens = len(tokenizer.encode(generated_part))
print(f"Generated {new_tokens} new tokens")
except KeyboardInterrupt:
print("\n\nInterrupted. Type /exit to quit.")
except Exception as e:
print(f"\nError: {e}")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="VDrontMoE Text Generation")
parser.add_argument("--prompt", type=str, help="Single prompt mode (no chat)")
parser.add_argument("--temperature", type=float, default=0.4, help="Temperature (default: 0.4)")
parser.add_argument("--max_tokens", type=int, default=100, help="Max new tokens")
parser.add_argument("--top_p", type=float, default=0.9, help="Nucleus sampling threshold")
parser.add_argument("--top_k", type=int, default=50, help="Top-K filtering")
parser.add_argument("--model_dir", type=str, default="./VDrontMoE-2m-5e")
args = parser.parse_args()
# Update parameters
MODEL_DIR = args.model_dir
TEMPERATURE = args.temperature
MAX_NEW_TOKENS = args.max_tokens
TOP_P = args.top_p
TOP_K = args.top_k
if args.prompt:
# Single mode
print(f"\nPrompt: {args.prompt}")
print(f"Temperature: {TEMPERATURE} | Max tokens: {MAX_NEW_TOKENS}")
print("\n" + "=" * 60)
result = generate(args.prompt)
print(result)
print("=" * 60)
else:
# Interactive mode
chat() |