from argparse import ArgumentParser from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig import deepspeed import torch from utils import DSPipeline
inputs = [ "DeepSpeed is a machine learning framework", "He is working on", "He has a", "He got all", "Everyone is happy and I can", "The new movie that got Oscar this year", "In the far far distance from our galaxy,", "Peace is the only way" ]
# Sample prompts. prompts = [ "Hello, my name is", "The president of the United States is", "The capital of France is", "The future of AI is", ] # Create a sampling params object. sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
# Create an LLM. llm = LLM(model="facebook/opt-125m") # Generate texts from the prompts. The output is a list of RequestOutput objects # that contain the prompt, generated text, and other information. outputs = llm.generate(prompts, sampling_params) # Print the outputs. for output in outputs: prompt = output.prompt generated_text = output.outputs[0].text print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")