Upload app.py

load model in 4bit for faster infer.

app.py

@@ -31,7 +31,7 @@ def load_model(model_name):
     model = AutoModelForCausalLM.from_pretrained(
         f"{root_path}/{model_name}",
         device_map="auto",
-        load_in_8bit=True,
+        load_in_4bit=True,
         torch_dtype=torch.bfloat16,
         trust_remote_code=True,
     )
@@ -47,7 +47,7 @@ with st.sidebar:
     top_p = st.sidebar.slider('top_p', min_value=0.01, max_value=1.0, value=0.9, step=0.01)
     top_k = st.sidebar.slider('top_k', min_value=1, max_value=1000, value=50, step=1)
     repetition_penalty = st.sidebar.slider('repetition penalty', min_value=1., max_value=2., value=1.2, step=0.05)
-    max_length = st.sidebar.slider('max new tokens', min_value=32, max_value=2000, value=240, step=8)
+    max_length = st.sidebar.slider('max new tokens', min_value=32, max_value=2000, value=512, step=8)
 
 with st.spinner('loading model..'):
     model, tokenizer = load_model(model_name)