File size: 2,851 Bytes
ccda2ec |
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 |
#!/usr/bin/env python
"""
Example usage script to evaluate a fine-tuned OlmoE adapter model
and demonstrate generation with adapters.
"""
import argparse
import torch
from transformers import AutoTokenizer
from modeling_olmoe import OlmoEWithAdaptersForCausalLM, OlmoConfig
def generate_text(
model_path: str,
prompt: str,
max_new_tokens: int = 128,
temperature: float = 0.7,
top_p: float = 0.9,
device: str = "auto",
):
"""Generate text using a fine-tuned OlmoE adapter model."""
# Determine device
if device == "auto":
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")
# Load tokenizer and model
print(f"Loading model from {model_path}")
tokenizer = AutoTokenizer.from_pretrained(model_path)
# Load config and update with adapter settings if needed
config = OlmoConfig.from_pretrained(model_path)
# Load adapter model
model = OlmoEWithAdaptersForCausalLM.from_pretrained(
model_path,
torch_dtype=torch.float16 if device == "cuda" else torch.float32,
)
model = model.to(device)
model.eval()
# Tokenize input
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
# Generate text
print("\nGenerating text...\n")
with torch.no_grad():
outputs = model.generate(
input_ids,
max_new_tokens=max_new_tokens,
do_sample=True,
temperature=temperature,
top_p=top_p,
)
# Decode the generated text
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(f"Prompt: {prompt}")
print("\nGenerated text:")
print("=" * 40)
print(generated_text)
print("=" * 40)
return generated_text
def main():
parser = argparse.ArgumentParser(description="Generate text with OlmoE adapter model")
parser.add_argument("--model_path", type=str, required=True, help="Path to the fine-tuned model")
parser.add_argument("--prompt", type=str, default="This is an example of", help="Prompt for text generation")
parser.add_argument("--max_new_tokens", type=int, default=128, help="Maximum number of new tokens to generate")
parser.add_argument("--temperature", type=float, default=0.7, help="Sampling temperature")
parser.add_argument("--top_p", type=float, default=0.9, help="Top-p sampling parameter")
parser.add_argument("--device", type=str, default="auto", help="Device to use (cuda, cpu, or auto)")
args = parser.parse_args()
generate_text(
model_path=args.model_path,
prompt=args.prompt,
max_new_tokens=args.max_new_tokens,
temperature=args.temperature,
top_p=args.top_p,
device=args.device,
)
if __name__ == "__main__":
main() |