Update app.py

app.py

@@ -1,26 +1,26 @@
 import spaces
 import os
-import subprocess
-import shlex
-if os.getenv('SYSTEM') == 'spaces':
-    git_repo = "https://github.com/huggingface/transformers.git"
-    subprocess.call(shlex.split(f'pip install git+{git_repo}'))
+# import subprocess
+# import shlex
+# if os.getenv('SYSTEM') == 'spaces':
+#     git_repo = "https://github.com/huggingface/transformers.git"
+#     subprocess.call(shlex.split(f'pip install git+{git_repo}'))
     
 import time
 import torch
-from transformers import OlmoeForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import gradio as gr
 from threading import Thread
 
-MODEL_LIST = ["allenai/OLMoE-1B-7B-0924-Instruct"]
+MODEL_LIST = ["openbmb/MiniCPM3-4B"]
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 MODEL = os.environ.get("MODEL_ID")
 
-TITLE = "<h1><center>OLMoE</center></h1>"
+TITLE = "<h1><center>MiniCPM3-4B</center></h1>"
 
 PLACEHOLDER = """
 <center>
-<p>Fully open, state-of-the-art Mixture of Expert model with 1.3 billion active and 6.9 billion total parameters.</p>
+<p>MiniCPM3-4B is the 3rd generation of MiniCPM series.</p>
 </center>
 """
 
@@ -39,19 +39,19 @@ h3 {
 
 device = "cuda" # for GPU usage or "cpu" for CPU usage
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL)
-model = OlmoeForCausalLM.from_pretrained(
+tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
     MODEL,
     torch_dtype=torch.bfloat16,
     device_map="auto",
-    ignore_mismatched_sizes=True)
+    trust_remote_code=True)
 
 @spaces.GPU()
 def stream_chat(
     message: str, 
     history: list, 
-    temperature: float = 0.5, 
-    max_new_tokens: int = 256, 
+    temperature: float = 0.7, 
+    max_new_tokens: int = 1024, 
     top_p: float = 1.0, 
     top_k: int = 20, 
     penalty: float = 1.2,
@@ -110,15 +110,15 @@ with gr.Blocks(css=CSS, theme="Nymbo/Nymbo_Theme") as demo:
                 minimum=0,
                 maximum=1,
                 step=0.1,
-                value=0.5,
+                value=0.7,
                 label="Temperature",
                 render=False,
             ),
             gr.Slider(
                 minimum=128,
-                maximum=4096,
+                maximum=32768,
                 step=1,
-                value=256,
+                value=1024,
                 label="Max new tokens",
                 render=False,
             ),