Update app.py

app.py

@@ -17,27 +17,33 @@ HF_TOKEN = st.secrets["HF_TOKEN"]
 st.set_page_config(page_title="DigiTwin RAG", page_icon="📂", layout="centered")
 st.title("📂 DigiTs the Twin")
 
-# --- File Upload UI ---
+# --- Sidebar ---
 with st.sidebar:
     st.header("📄 Upload Knowledge Files")
     uploaded_files = st.file_uploader("Upload PDFs or .txt files", accept_multiple_files=True, type=["pdf", "txt"])
+    model_choice = st.selectbox("🧠 Choose Model", ["Qwen", "Mistral"])
     if uploaded_files:
         st.success(f"{len(uploaded_files)} file(s) uploaded")
 
 # --- Load Model & Tokenizer ---
 @st.cache_resource
-def load_model():  #  amiguel/GM_Qwen1.8B_Finetune
-    tokenizer = AutoTokenizer.from_pretrained("amiguel/GM_Mistral7B_Finetune", trust_remote_code=True, token=HF_TOKEN)
+def load_model(selected_model):
+    if selected_model == "Qwen":
+        model_id = "amiguel/GM_Qwen1.8B_Finetune"
+    else:
+        model_id = "amiguel/GM_Mistral7B_Finetune"
+
+    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True, token=HF_TOKEN)
     model = AutoModelForCausalLM.from_pretrained(
-        "amiguel/GM_Mistral7B_Finetune",
+        model_id,
         device_map="auto",
         torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32,
         trust_remote_code=True,
         token=HF_TOKEN
     )
-    return model, tokenizer
+    return model, tokenizer, model_id
 
-model, tokenizer = load_model()
+model, tokenizer, model_id = load_model(model_choice)
 
 # --- System Prompt ---
 SYSTEM_PROMPT = (
@@ -128,14 +134,19 @@ if prompt := st.chat_input("Ask something based on uploaded documents..."):
 
         for chunk in generate_response(full_prompt):
             answer += chunk
-            container.markdown(answer + "▌", unsafe_allow_html=True)
-        container.markdown(answer)
+            cleaned = answer
+
+            # 🔧 Strip <|im_start|>, <|im_end|> if using Mistral (Qwen needs them)
+            if "Mistral" in model_id:
+                cleaned = cleaned.replace("<|im_start|>", "").replace("<|im_end|>", "").strip()
+
+            container.markdown(cleaned + "▌", unsafe_allow_html=True)
 
         end = time.time()
-        st.session_state.messages.append({"role": "assistant", "content": answer})
+        st.session_state.messages.append({"role": "assistant", "content": cleaned})
 
         input_tokens = len(tokenizer(full_prompt)["input_ids"])
-        output_tokens = len(tokenizer(answer)["input_ids"])
+        output_tokens = len(tokenizer(cleaned)["input_ids"])
         speed = output_tokens / (end - start)
 
         with st.expander("📊 Debug Info"):