add mistral code

app.py

@@ -1,4 +1,40 @@
 import streamlit as st
 
 x = st.slider('Select a value')
-st.write(x, 'squared is', x * x)
+st.write(x, 'squared is', x * x)
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# Set the device to CUDA if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+print("Running on GPU: ", torch.cuda.is_available())
+
+model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1", torch_dtype="auto", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1", trust_remote_code=True)
+
+# Set the padding token if not already defined
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+# Move model to the selected device
+model = model.to(device)
+
+#input_text = "What are the side effects of sunscreen?"
+
+input_text = st.text_input()
+
+if st.button("generate response"):
+    # Encode input text along with attention mask
+    encoding = tokenizer(input_text, return_tensors='pt', max_length=150, padding='max_length') 
+    attention_mask = encoding['attention_mask'].to(device)
+
+    # Move input tensors to the same device as the model
+    inputs = encoding['input_ids'].to(device)
+
+    # Generate output using both input_ids and attention_mask
+    outputs = model.generate(inputs, attention_mask=attention_mask, max_new_tokens= 50, num_return_sequences=1)
+
+    for i, output_id in enumerate(outputs):
+        st.write(f"Generated text {i+1}: {tokenizer.decode(output_id, skip_special_tokens=True)}")