Update main.py

main.py

@@ -1,20 +1,9 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from mistral_inference.transformer import Transformer
+from mistral_inference.generate import generate
 
-# Load model and tokenizer
-model_name = "mistralai/Codestral-22B-v0.1"
-model = AutoModelForCausalLM.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = Transformer.from_folder(mistral_models_path)
+out_tokens, _ = generate([tokens], model, max_tokens=64, temperature=0.0, eos_id=tokenizer.instruct_tokenizer.tokenizer.eos_id)
 
-# No need to move model to GPU, default is CPU
-# model.to("cpu")  # This line can be omitted since it's already on CPU by default
+result = tokenizer.decode(out_tokens[0])
 
-# Encode input tokens
-input_text = "Your input text here"
-tokens = tokenizer(input_text, return_tensors="pt").input_ids
-
-# Generate output
-generated_ids = model.generate(tokens, max_new_tokens=1000, do_sample=True)
-
-# Decode generated tokens
-result = tokenizer.decode(generated_ids[0].tolist(), skip_special_tokens=True)
 print(result)