Update app.py

app.py

@@ -1,36 +1,33 @@
 import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Load the Biomistral 7b model and tokenizer
-model_name = "biomistral/Biomistral-7b"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
+@st.cache_resource
+def load_model_and_tokenizer():
+    model_name_or_path = "m42-health/med42-70b"
+    model = AutoModelForCausalLM.from_pretrained(model_name_or_path, device_map="auto")
+    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
+    return model, tokenizer
 
-# Define the text generation function
-def generate_text(prompt, max_length=500, num_return_sequences=1, temperature=0.7):
-    input_ids = tokenizer.encode(prompt, return_tensors="pt")
-    output = model.generate(
-        input_ids,
-        max_length=max_length,
-        num_return_sequences=num_return_sequences,
-        temperature=temperature,
-        pad_token_id=tokenizer.eos_token_id,
-    )
-    generated_text = tokenizer.batch_decode(output, skip_special_tokens=True)
-    return generated_text
+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)
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
+    return response
 
-# Streamlit app
 def main():
-    st.title("Doctor Chatbot (Powered by Biomistral 7b)")
-    st.write("Welcome to the Doctor Chatbot. Please describe your symptoms or ask a medical question, and I'll provide a response.")
+    st.title("M42 Health Medical Assistant")
+    model, tokenizer = load_model_and_tokenizer()
 
-    user_input = st.text_area("Enter your symptoms or question:")
-
-    if user_input:
+    prompt = st.text_area("Enter your medical query:")
+    if st.button("Submit"):
         with st.spinner("Generating response..."):
-            generated_text = generate_text(user_input)
-            st.write(generated_text[0])
+            response = generate_response(prompt)
+            st.write(response)
 
 if __name__ == "__main__":
-    main()
-
+    main()