Update app.py

app.py

@@ -8,7 +8,7 @@ import spaces
 
 import torch
 import gradio as gr
-from transformers import Mistral3ForConditionalGeneration, AutoProcessor
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 from dd import encode_prompt
 
@@ -34,13 +34,13 @@ DTYPE = torch.bfloat16
 logger.info("Loading models...")
 
 t0 = time.time()
-text_encoder = Mistral3ForConditionalGeneration.from_pretrained(
+text_encoder = AutoModelForCausalLM.from_pretrained(
     TEXT_ENCODER_ID,
     dtype=DTYPE,
 ).to("cpu")
 
 t1 = time.time()
-tokenizer = AutoProcessor.from_pretrained(TOKENIZER_ID)
+tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_ID)
 logger.info("Loaded tokenizer in %.2fs", time.time() - t1)
 
 torch.set_grad_enabled(False)
@@ -132,7 +132,7 @@ with gr.Blocks(title="Mistral Text Encoder") as demo:
 
     encode_btn.click(
         fn=encode_text,
-        inputs=[prompt_input],
+        inputs=[{"role": "user",prompt_input}],
         outputs=[output_file, status_output],
     )