Update app.py

app.py

@@ -1,14 +1,34 @@
 # Import necessary libraries
 import streamlit as st
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoModelForCausalLM
+from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
+from mistral_common.protocol.instruct.messages import UserMessage
+from mistral_common.protocol.instruct.request import ChatCompletionRequest
+import torch
 
-# Load the model and tokenizer
-model_name = "mistralai/Codestral-22B-v0.1"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+# Path to the mistral models
+mistral_models_path = "MISTRAL_MODELS_PATH"
 
-# Initialize the pipeline
-text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+# Load the tokenizer
+tokenizer = MistralTokenizer.v3()
+
+# Load the model
+model = AutoModelForCausalLM.from_pretrained("mistralai/Codestral-22B-v0.1")
+model.to("cuda")
+
+# Function to generate text
+def generate_text(prompt):
+    # Encode the prompt
+    completion_request = ChatCompletionRequest(messages=[UserMessage(content=prompt)])
+    tokens = tokenizer.encode_chat_completion(completion_request).tokens
+
+    # Generate text using the model
+    with torch.no_grad():
+        generated_ids = model.generate(torch.tensor([tokens]).to(model.device), max_new_tokens=1000, do_sample=True)
+
+    # Decode the generated text
+    result = tokenizer.decode(generated_ids[0].tolist())
+    return result
 
 # Streamlit interface
 st.title("Codestral Text Generation")
@@ -25,9 +45,9 @@ if st.button("Generate"):
     if user_input:
         with st.spinner("Generating text..."):
             # Generate text using the model
-            generated_text = text_generator(user_input, max_length=100, num_return_sequences=1)
+            generated_text = generate_text(user_input)
             st.write("### Generated Text")
-            st.write(generated_text[0]['generated_text'])
+            st.write(generated_text)
     else:
         st.warning("Please enter a prompt to generate text.")