Spaces:
Paused
Paused
| from transformers import AutoTokenizer, pipeline | |
| import torch | |
| import streamlit as st | |
| def load(): | |
| model = "meta-llama/Llama-2-13b-chat-hf" | |
| tokenizer = AutoTokenizer.from_pretrained(model) | |
| pipeline_generator = pipeline( | |
| "text-generation", | |
| model=model, | |
| torch_dtype=torch.float16, | |
| device_map="auto", | |
| ) | |
| st.session_state.model = model | |
| st.session_state.tokenizer = tokenizer | |
| st.session_state.pipeline_generator = pipeline_generator | |
| # load() |