Update app.py

app.py

@@ -1,13 +1,13 @@
 import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import subprocess
-subprocess.run(["pip", "install", "accelerate"])
+import accelerate
 
 # Load the model and tokenizer
 @st.cache_resource
 def load_model_and_tokenizer():
-    model_name_or_path = "mistralai/Mistral-7B-Instruct-v0.2"
-    model = AutoModelForCausalLM.from_pretrained(model_name_or_path, device_map="auto")
+    model_name_or_path = "anthropic/mistral-7b"
+    accelerator = accelerate.Accelerator(device_map="auto")
+    model = AutoModelForCausalLM.from_pretrained(model_name_or_path, device_map=accelerator.device_map)
     tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
     return model, tokenizer
 
@@ -15,21 +15,21 @@ def load_model_and_tokenizer():
 @st.cache_data
 def generate_response(prompt):
     prompt_template = f'''
-    <|system|>: You are a helpful medical assistant created by M42 Health in the UAE.
     <|prompter|>:{prompt}
     <|assistant|>:
     '''
-    input_ids = tokenizer(prompt_template, return_tensors='pt').input_ids.cuda()
-    output = model.generate(inputs=input_ids, temperature=0.7, do_sample=True, eos_token_id=tokenizer.eos_token_id, pad_token_id=tokenizer.pad_token_id, max_new_tokens=512)
+    input_ids = tokenizer(prompt_template, return_tensors='pt').input_ids
+    with accelerator.autocast():
+        output = model.generate(inputs=input_ids, temperature=0.7, do_sample=True, eos_token_id=tokenizer.eos_token_id, pad_token_id=tokenizer.pad_token_id, max_new_tokens=512)
     response = tokenizer.decode(output[0], skip_special_tokens=True)
     return response
 
 # Streamlit app
 def main():
-    st.title("Med42 - Clinical Large Language Model")
+    st.title("Mistral 7B Language Model")
     model, tokenizer = load_model_and_tokenizer()
 
-    prompt = st.text_area("Enter your medical query:")
+    prompt = st.text_area("Enter your query:")
     if st.button("Submit"):
         with st.spinner("Generating response..."):
             response = generate_response(prompt)