Update app.py

app.py

@@ -8,29 +8,43 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
 import torch
 
 # ──────────────────────────────────────────────
-MODEL_REPO = "bugaku/fufufafa"
+MODELS = {
+    "3B": "bugaku/igpt",
+    "1.5B": "bugaku/igpt-smol",
+}
 # ──────────────────────────────────────────────
 
-print(f"Loading model: {MODEL_REPO}")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
+loaded_models = {}
 
-config = AutoConfig.from_pretrained(MODEL_REPO)
-if hasattr(config, "quantization_config"):
-    delattr(config, "quantization_config")
-if hasattr(config, "_pre_quantization_dtype"):
-    delattr(config, "_pre_quantization_dtype")
+def load_model(model_repo):
+    if model_repo in loaded_models:
+        return loaded_models[model_repo]
 
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_REPO,
-    config=config,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    device_map="auto",
-)
-model.eval()
-print("Model loaded!")
+    print(f"Loading model: {model_repo}")
+    tokenizer = AutoTokenizer.from_pretrained(model_repo)
 
+    config = AutoConfig.from_pretrained(model_repo)
+    if hasattr(config, "quantization_config"):
+        delattr(config, "quantization_config")
+    if hasattr(config, "_pre_quantization_dtype"):
+        delattr(config, "_pre_quantization_dtype")
+
+    model = AutoModelForCausalLM.from_pretrained(
+        model_repo,
+        config=config,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        device_map="auto",
+    )
+    model.eval()
+    loaded_models[model_repo] = (model, tokenizer)
+    print(f"Loaded: {model_repo}")
+    return model, tokenizer
+
+
+def respond(message, history, model_choice, max_new_tokens, temperature):
+    model_repo = MODELS[model_choice]
+    model, tokenizer = load_model(model_repo)
 
-def respond(message, history, max_new_tokens, temperature):
     messages = [{"role": "user", "content": message}]
 
     encoded = tokenizer.apply_chat_template(
@@ -62,10 +76,11 @@ with gr.Blocks(title="Style Model") as demo:
     chatbot = gr.ChatInterface(
         fn=respond,
         additional_inputs=[
+            gr.Dropdown(choices=list(MODELS.keys()), value=list(MODELS.keys())[0], label="Model"),
             gr.Slider(32, 512, value=256, step=32, label="Max response length"),
             gr.Slider(0.1, 2.0, value=0.8, step=0.1, label="Temperature"),
         ],
         chatbot=gr.Chatbot(height=500),
     )
 
-demo.launch()
+demo.launch()