Update app.py

app.py

@@ -21,7 +21,7 @@ st.title("📂 DigiTs the Twin")
 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"])
+    model_choice = st.selectbox("🧠 Choose Model", ["Qwen", "Mistral", "Llama3"])
     if uploaded_files:
         st.success(f"{len(uploaded_files)} file(s) uploaded")
 
@@ -30,10 +30,8 @@ with st.sidebar:
 def load_model(selected_model):
     if selected_model == "Qwen":
         model_id = "amiguel/GM_Qwen1.8B_Finetune"
-
-    elif selected_model == "Llama":
-          model_id = "amiguel/Llama3_8B_Instruct_FP16"
-    
+    elif selected_model == "Llama3":
+        model_id = "amiguel/Llama3_8B_Instruct_FP16"
     else:
         model_id = "amiguel/GM_Mistral7B_Finetune"
 
@@ -61,7 +59,6 @@ SYSTEM_PROMPT = (
 # --- Prompt Builder ---
 def build_prompt(messages, context="", model_name="Qwen"):
     if "Mistral" in model_name:
-        # Alpaca-style prompt
         prompt = f"You are DigiTwin, an expert in offshore inspection, maintenance, and asset integrity.\n"
         if context:
             prompt += f"Here is relevant context:\n{context}\n\n"
@@ -71,8 +68,18 @@ def build_prompt(messages, context="", model_name="Qwen"):
             elif msg["role"] == "assistant":
                 prompt += f"### Response:\n{msg['content'].strip()}\n"
         prompt += "### Response:\n"
-    else:
-        # Qwen-style
+
+    elif "Llama" in model_name:
+        prompt = "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n"
+        prompt += f"{SYSTEM_PROMPT}\n\nContext:\n{context}\n"
+        for msg in messages:
+            if msg["role"] == "user":
+                prompt += "<|start_header_id|>user<|end_header_id|>\n" + msg["content"].strip() + "\n"
+            elif msg["role"] == "assistant":
+                prompt += "<|start_header_id|>assistant<|end_header_id|>\n" + msg["content"].strip() + "\n"
+        prompt += "<|start_header_id|>assistant<|end_header_id|>\n"
+
+    else:  # Qwen
         prompt = f"<|im_start|>system\n{SYSTEM_PROMPT}\n\nContext:\n{context}<|im_end|>\n"
         for msg in messages:
             role = msg["role"]
@@ -80,7 +87,6 @@ def build_prompt(messages, context="", model_name="Qwen"):
         prompt += "<|im_start|>assistant\n"
     return prompt
 
-
 # --- Embed Uploaded Documents ---
 @st.cache_resource
 def embed_uploaded_files(files):
@@ -125,7 +131,7 @@ BOT_AVATAR = "https://raw.githubusercontent.com/achilela/vila_fofoka_analysis/99
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-# --- Display Message History ---
+# --- Display Chat History ---
 for msg in st.session_state.messages:
     with st.chat_message(msg["role"], avatar=USER_AVATAR if msg["role"] == "user" else BOT_AVATAR):
         st.markdown(msg["content"])
@@ -141,7 +147,6 @@ if prompt := st.chat_input("Ask something based on uploaded documents..."):
         docs = retriever.similarity_search(prompt, k=3)
         context = "\n\n".join([doc.page_content for doc in docs])
 
-    # Limit to last 6 messages for memory
     recent_messages = st.session_state.messages[-6:]
     full_prompt = build_prompt(recent_messages, context, model_name=model_id)
 
@@ -154,9 +159,10 @@ if prompt := st.chat_input("Ask something based on uploaded documents..."):
             answer += chunk
             cleaned = answer
 
-            # 🔧 Strip <|im_start|>, <|im_end|> if using Mistral (Qwen needs them)
-            if "Mistral" in model_id:
+            if "Mistral" in model_id or "Llama" in model_id:
                 cleaned = cleaned.replace("<|im_start|>", "").replace("<|im_end|>", "").strip()
+                cleaned = cleaned.replace("<|start_header_id|>", "").replace("<|end_header_id|>", "")
+                cleaned = cleaned.replace("<|begin_of_text|>", "").strip()
 
             container.markdown(cleaned + "▌", unsafe_allow_html=True)