Update app.py

app.py

@@ -2,19 +2,16 @@ import os
 import time
 import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import gradio as gr
 from threading import Thread
 
-MODEL_LIST = ["mistralai/Ministral-8B-Instruct-2410"]
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
-MODEL = os.environ.get("MODEL_ID")
 
-TITLE = "<h1><center>Mistral-Nemo</center></h1>"
+TITLE = "<h1><center>Mistral-lab</center></h1>"
 
 PLACEHOLDER = """
 <center>
-<p>The Mistral-8B is a pretrained generative text model of 8B parameters trained jointly by Mistral AI.</p>
+<p>Chat with Mistral AI LLM.</p>
 </center>
 """
 
@@ -31,14 +28,26 @@ h3 {
 }
 """
 
+from huggingface_hub import snapshot_download
+from pathlib import Path
+
+mistral_models_path = Path.home().joinpath('mistral_models', '8B-Instruct')
+mistral_models_path.mkdir(parents=True, exist_ok=True)
+
+snapshot_download(repo_id="mistralai/Ministral-8B-Instruct-2410", allow_patterns=["params.json", "consolidated.safetensors", "tekken.json"], local_dir=mistral_models_path)
+
+from mistral_inference.transformer import Transformer
+from mistral_inference.generate import generate
+
+from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
+from mistral_common.protocol.instruct.messages import AssistantMessage, UserMessage
+from mistral_common.protocol.instruct.request import ChatCompletionRequest
+
 device = "cuda" # for GPU usage or "cpu" for CPU usage
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL,
-    torch_dtype=torch.bfloat16,
-    device_map="auto",
-    ignore_mismatched_sizes=True)
+tokenizer = MistralTokenizer.from_file(f"{mistral_models_path}/tekken.json")
+model = Transformer.from_folder(mistral_models_path)
+
 
 @spaces.GPU()
 def stream_chat(
@@ -55,42 +64,31 @@ def stream_chat(
 
     conversation = []
     for prompt, answer in history:
-        conversation.extend([
-            {"role": "user", "content": prompt}, 
-            {"role": "assistant", "content": answer},
-        ])
-
-    conversation.append({"role": "user", "content": message})
+        conversation.append(UserMessage(content=prompt))
+        conversation.append(AssistantMessage(content=answer))
+    conversation.append(UserMessage(content=message))
 
-    input_text=tokenizer.apply_chat_template(conversation, tokenize=False)
-    inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
-    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
+    completion_request = ChatCompletionRequest(messages=conversation)
     
-    generate_kwargs = dict(
-        input_ids=inputs, 
-        max_new_tokens = max_new_tokens,
-        do_sample = False if temperature == 0 else True,
+    tokens = tokenizer.encode_chat_completion(completion_request).tokens
+    
+    out_tokens, _ = generate(
+        [tokens], 
+        model, 
+        max_tokens=max_new_tokens, 
+        temperature=temperature,
         top_p = top_p,
         top_k = top_k,
-        temperature = temperature,
-        streamer=streamer,
         repetition_penalty=penalty,
-        pad_token_id = 10,
-    )
-
-    with torch.no_grad():
-        thread = Thread(target=model.generate, kwargs=generate_kwargs)
-        thread.start()
-        
-    buffer = ""
-    for new_text in streamer:
-        buffer += new_text
-        yield buffer
-
+        eos_id=tokenizer.instruct_tokenizer.tokenizer.eos_id)
+    
+    result = tokenizer.instruct_tokenizer.tokenizer.decode(out_tokens[0])
+    
+    return result
             
 chatbot = gr.Chatbot(height=600, placeholder=PLACEHOLDER)
 
-with gr.Blocks(css=CSS, theme="Nymbo/Nymbo_Theme") as demo:
+with gr.Blocks(css=CSS, theme="ocean") as demo:
     gr.HTML(TITLE)
     gr.DuplicateButton(value="Duplicate Space for private use", elem_classes="duplicate-button")
     gr.ChatInterface(