Update app.py

app.py

@@ -7,7 +7,6 @@ import torch
 from peft import LoraConfig, PeftModel, prepare_model_for_kbit_training, get_peft_model
 from huggingface_hub import login
 
-
 # Initialize the OpenAI client (if needed for Hugging Face Inference API)
 client = OpenAI(
     base_url="https://api-inference.huggingface.co/v1",
@@ -15,7 +14,6 @@ client = OpenAI(
 )
 
 api_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
-
 if api_token:
     login(token=api_token)
 else:
@@ -23,15 +21,11 @@ else:
 
 # Define model links and configurations
 model_links = {
-    "HAH-2024-v0.1": "drmasad/HAH-2024-v0.11",
-    "Mistral": "mistralai/Mistral-7B-Instruct-v0.2",
+    "HAH-2024-v0.1": "drmasad/HAH-2024-v0.11"
 }
 
 # Define sidebar options
-models = list(model_links.keys())
-
-# Sidebar model selection
-selected_model = st.sidebar.selectbox("Select Model", models)
+selected_model = "HAH-2024-v0.1"  # Directly using your model
 
 # Sidebar temperature control
 temp_values = st.sidebar.slider("Select a temperature value", 0.0, 1.0, (0.5))
@@ -48,102 +42,80 @@ model_info = {
     "HAH-2024-v0.1": {
         "description": "HAH-2024-v0.1 is a fine-tuned model based on Mistral 7B. It's designed for conversations on diabetes.",
         "logo": "https://www.hmgaihub.com/untitled.png",
-    },
-    "Mistral": {
-        "description": "Mistral is a large language model with multi-task capabilities.",
-        "logo": "https://mistral.ai/images/logo_hubc88c4ece131b91c7cb753f40e9e1cc5_2589_256x0_resize_q97_h2_lanczos_3.webp",
-    },
+    }
 }
 
 st.sidebar.write(f"You're now chatting with **{selected_model}**")
 st.sidebar.markdown(model_info[selected_model]["description"])
 st.sidebar.image(model_info[selected_model]["logo"])
 
-# Load the appropriate model based on user selection
-def load_model(selected_model_name):
-    if selected_model_name == "HAH-2024-v0.1":
-        # Setup for HAH-2024-v0.1
-        model_name = model_links["HAH-2024-v0.1"]
-        base_model = "mistralai/Mistral-7B-Instruct-v0.2"
-
-        # Load model with quantization configuration
-        bnb_config = BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_quant_type="nf4",
-            bnb_4bit_compute_dtype=torch.bfloat16,
-            bnb_4bit_use_double_quant=False,
-        )
-
-        model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            quantization_config=bnb_config,
-            torch_dtype=torch.bfloat16,
-            device_map="auto",
-            trust_remote_code=True,
-        )
-
-        model.config.use_cache = False
-        model = prepare_model_for_kbit_training(model)
-
-        peft_config = LoraConfig(
-            lora_alpha=16,
-            lora_dropout=0.1,
-            r=64,
-            bias="none",
-            task_type="CAUSAL_LM",
-            target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj"],
-        )
-
-        model = get_peft_model(model, peft_config)
-
-        tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)
-
-    elif selected_model_name == "Mistral":
-        # Setup for Mistral 7B
-        model = AutoModelForCausalLM.from_pretrained(
-            model_links[selected_model_name]
-        )
-        tokenizer = AutoTokenizer.from_pretrained(model_links[selected_model_name])
+# Load the appropriate model
+def load_model():
+    model_name = model_links["HAH-2024-v0.1"]
+    base_model = "mistralai/Mistral-7B-Instruct-v0.2"
+
+    # Load model with quantization configuration
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.bfloat16,
+        bnb_4bit_use_double_quant=False,
+    )
+
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        quantization_config=bnb_config,
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+        trust_remote_code=True,
+    )
+
+    model.config.use_cache = False
+    model = prepare_model_for_kbit_training(model)
+
+    peft_config = LoraConfig(
+        lora_alpha=16,
+        lora_dropout=0.1,
+        r=64,
+        bias="none",
+        task_type="CAUSAL_LM",
+        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj"],
+    )
+
+    model = get_peft_model(model, peft_config)
+
+    tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)
 
     return model, tokenizer
 
+model, tokenizer = load_model()
+
 # Initialize chat history
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-# Load the selected model
-model, tokenizer = load_model(selected_model)
-
-st.subheader(f"AI - {selected_model}")
-
 # Display previous chat messages
 for message in st.session_state.messages:
     with st.chat_message(message["role"]):
         st.markdown(message["content"])
 
 # User input for conversation
-if prompt := st.chat_input("Ask a question"):
-    # Display user input
+if prompt := st.chat_input("Ask me anything about diabetes"):
     with st.chat_message("user"):
         st.markdown(prompt)
-
-    # Store the user message
+    
     st.session_state.messages.append({"role": "user", "content": prompt})
-
-    # Generate the assistant's response
+    
     with st.chat_message("assistant"):
-        pipe = pipeline(
+        result = pipeline(
             task="text-generation", 
             model=model, 
             tokenizer=tokenizer, 
             max_length=1024, 
             temperature=temp_values
-        )
+        )(prompt)
         
-        result = pipe(f"<s>[INST] {prompt}</s>", do_sample=True)
-        response = result[0]["generated_text"]
-
+        response = result[0]['generated_text']
         st.markdown(response)
-
-    # Store the assistant's response
+    
     st.session_state.messages.append({"role": "assistant", "content": response})