Update app.py

app.py

@@ -5,17 +5,18 @@ from dotenv import load_dotenv
 from xhtml2pdf import pisa
 import io
 from textwrap import dedent
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline  # for loading llama tokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
-# --- Load Model Resources ---
+# Load Model Resources
 def load_resources():
     load_dotenv()
     huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
     subprocess.run(["huggingface-cli", "login", "--token", huggingface_token], capture_output=True)
-    tokenizer = AutoTokenizer.from_pretrained("istiak101/TinyLlama-1.1B-Chat-v0.6-rag-finetunedv1.0")
-    model = AutoModelForCausalLM.from_pretrained("istiak101/TinyLlama-1.1B-Chat-v0.6-rag-finetunedv1.0")
+    tokenizer = AutoTokenizer.from_pretrained("istiak101/TinyLlama-1.1B-Chat-v0.6-rag-finetunedv2.0")
+    model = AutoModelForCausalLM.from_pretrained("istiak101/TinyLlama-1.1B-Chat-v0.6-rag-finetunedv2.0")
     return model, tokenizer
 
+# Chat Prompt
 def create_test_prompt(question, context, tokenizer):
     prompt = dedent(
         f"""
@@ -39,7 +40,7 @@ def create_test_prompt(question, context, tokenizer):
         messages, tokenize=False, add_generation_prompt=True
     )
 
-# --- Store model and tokenizer in session state ---
+# Store model and tokenizer in session state
 if "llama_model" not in st.session_state or "llama_tokenizer" not in st.session_state:
     model, tokenizer = load_resources()
     st.session_state.llama_model = model
@@ -64,7 +65,7 @@ def get_llama_response(query):
     outputs = pipe(query)
     return outputs[0]["generated_text"]
 
-# --- PDF Generation ---
+# PDF Generation
 def generate_pdf(convo, topic):
     html = f"<h2>{topic}</h2><hr>"
     for msg in convo:
@@ -79,7 +80,7 @@ def generate_pdf(convo, topic):
         return None
     return result
 
-# --- Session Init ---
+# Session Init
 if "chat_sessions" not in st.session_state:
     st.session_state.chat_sessions = {}
 if "current_conversation" not in st.session_state:
@@ -87,10 +88,10 @@ if "current_conversation" not in st.session_state:
 if "edit_mode" not in st.session_state:
     st.session_state.edit_mode = {}
 
-# --- App Title ---
+# App Title
 st.title("💬 Tell Me Why")
 
-# --- Custom CSS ---
+# Custom CSS
 st.markdown("""
 <style>
 .chat-wrapper {
@@ -150,7 +151,7 @@ st.markdown("""
 </style>
 """, unsafe_allow_html=True)
 
-# --- Sidebar: Conversations ---
+# Sidebar: Conversations
 st.sidebar.title("Conversations")
 titles = list(st.session_state.chat_sessions.keys())
 
@@ -168,7 +169,7 @@ if titles:
 else:
     st.sidebar.write("No conversations yet. Start one below!")
 
-# --- New Conversation ---
+# New Conversation
 with st.sidebar.form(key='new_conversation_form', clear_on_submit=True):
     new_topic = st.text_input("New Conversation Name")
     submit_button = st.form_submit_button("Start New Conversation")
@@ -184,7 +185,7 @@ with st.sidebar.form(key='new_conversation_form', clear_on_submit=True):
         else:
             st.sidebar.warning("Conversation already exists!")
 
-# --- Main Chat Area ---
+# Main Chat Area
 if st.session_state.current_conversation:
     convo = st.session_state.chat_sessions[st.session_state.current_conversation]
 
@@ -249,8 +250,7 @@ if st.session_state.current_conversation:
 
     st.markdown('</div>', unsafe_allow_html=True)
 
-    # --- User Prompt ---
-    # --- Generate Assistant Response ---
+    # User Prompt
     if len(convo) % 2 == 1:
         last_user_msg = convo[-1]["text"]
         with st.spinner("Generating response..."):