v1

app.py

@@ -22,6 +22,9 @@ from torchvision.transforms.functional import pil_to_tensor
 # accel
 accel = Accelerator()
 
+# loading model
+model_0_5, tokenizer_0_5 = load_model(size='0.5b')
+
 # loading model
 model_1_8, tokenizer_1_8 = load_model(size='1.8b')
 
@@ -50,7 +53,10 @@ def threading_function(inputs, streamer, device, model, tokenizer, temperature,
 def bot_streaming(message, history, link, temperature, new_max_token, top_p):
 
     # model selection
-    if "1.8B" in link:
+    if "0.5B" in link:
+        model = model_0_5
+        tokenizer = tokenizer_0_5
+    elif "1.8B" in link:
         model = model_1_8
         tokenizer = tokenizer_1_8
     elif "3.8B" in link:
@@ -127,7 +133,7 @@ def bot_streaming(message, history, link, temperature, new_max_token, top_p):
         yield buffer
 
 demo = gr.ChatInterface(fn=bot_streaming,
-                        additional_inputs = [gr.Radio(["1.8B", "3.8B", "7B"], label="Size", info="Select one model size", value="7B"), gr.Slider(0, 1, 0.9, label="temperature"), gr.Slider(1, 1024, 128, label="new_max_token"), gr.Slider(0, 1, 0.95, label="top_p")],
+                        additional_inputs = [gr.Radio(["0.5B", "1.8B", "3.8B", "7B"], label="Size", info="Select one model size", value="7B"), gr.Slider(0, 1, 0.9, label="temperature"), gr.Slider(1, 1024, 128, label="new_max_token"), gr.Slider(0, 1, 0.95, label="top_p")],
                         additional_inputs_accordion="Generation Hyperparameters",
                         theme=gr.themes.Soft(),
                         title="Phantom",