Spaces:
Sleeping
Sleeping
File size: 1,242 Bytes
f6edbce c5d8529 f6edbce c5d8529 |
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 |
import streamlit as st
import transformers
import torch
import json
import os
from transformers import AutoTokenizer, TextStreamer , pipeline
model_id = "WizardLM/WizardMath-7B-V1.1"
# Configuration
runtimeFlag = "cuda:0" #Run on GPU (you can't run GPTQ on cpu)
cache_dir = None # by default, don't set a cache directory. This is automatically updated if you connect Google Drive.
scaling_factor = 1.0 # allows for a max sequence length of 16384*6 = 98304! Unfortunately, requires Colab Pro and a V100 or A100 to have sufficient RAM.
from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
model_id,
device_map="auto",
offload_folder="offload",
pad_token_id=tokenizer.eos_token_id,
offload_state_dict = True,
torch_dtype=torch.float16,
# rope_scaling = {"type": "dynamic", "factor": scaling_factor}
)
pipe = pipeline(
"text-generation",
model=model,
tokenizer=tokenizer,
max_new_tokens=512,
temperature=0.7,
top_p=0.95,
repetition_penalty=1.15
)
question = st.text_area("Enter questoin")
if text:
out = pipe(question)[0]['generated_text']
st.write(out) |