| from transformers import AutoTokenizer | |
| import transformers | |
| import torch | |
| model = "meta-llama/Llama-2-7b-chat-hf" | |
| tokenizer = AutoTokenizer.from_pretrained(model) | |
| pipeline = transformers.pipeline( | |
| "text-generation", | |
| model=model, | |
| torch_dtype=torch.float16, | |
| device_map="auto", | |
| ) | |
| sequences = pipeline( | |
| 'I liked "Breaking Bad" and "Band of Brothers". Do you have any recommendations of other shows I might like?\n', | |
| do_sample=True, | |
| top_k=10, | |
| num_return_sequences=1, | |
| eos_token_id=tokenizer.eos_token_id, | |
| max_length=200, | |
| ) | |
| for seq in sequences: | |
| print(f"Result: {seq['generated_text']}") |