Llama2_RAG / app.py
Waflon's picture
Create app.py
5be1431 verified
raw
history blame contribute delete
487 Bytes
from transformers import AutoTokenizer
import transformers
import torch
import accelerate
model = "meta-llama/Llama-2-7b-chat-hf"
tokenizer=AutoTokenizer.from_pretrained(model)
pipeline=transformers.pipeline(
"text-generation",
model=model,
tokenizer=tokenizer,
torch_dtype=torch.bfloat16,
trust_remote_code=True,
device_map="auto",
max_length=1000,
do_sample=True,
top_k=10,
num_return_sequences=1,
eos_token_id=tokenizer.eos_token_id
)