Update app.py

app.py

@@ -4,17 +4,41 @@ import streamlit as st
 from dotenv import load_dotenv
 from xhtml2pdf import pisa
 import io
-from transformers import AutoTokenizer, AutoModelForCausalLM  # for loading llama tokenizer
+from textwrap import dedent
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline  # for loading llama tokenizer
 
 # --- 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-Finetuned")
-    model = AutoModelForCausalLM.from_pretrained("istiak101/TinyLlama-1.1B-Finetuned")
+    tokenizer = AutoTokenizer.from_pretrained("istiak101/TinyLlama-1.1B-Finetunedv1.0")
+    model = AutoModelForCausalLM.from_pretrained("istiak101/TinyLlama-1.1B-Finetunedv1.0")
     return model, tokenizer
 
+def create_test_prompt(question, context, tokenizer):
+    prompt = dedent(
+        f"""
+    {question}
+
+    Information:
+
+    ```
+    {context}
+    ```
+    """
+    )
+    messages = [
+        {
+            "role": "system",
+            "content": "Use only the information to answer the question",
+        },
+        {"role": "user", "content": prompt},
+    ]
+    return tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+
 # --- 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()
@@ -24,13 +48,21 @@ if "llama_model" not in st.session_state or "llama_tokenizer" not in st.session_
 st.set_page_config(page_title="Ask Wikipedia", page_icon="📘", layout="wide")
 
 def get_llama_response(query):
-    model = st.session_state.llama_model
-    tokenizer = st.session_state.llama_tokenizer
+    # model = st.session_state.llama_model
+    # tokenizer = st.session_state.llama_tokenizer
 
-    inputs = tokenizer(query, return_tensors="pt")
-    outputs = model.generate(**inputs, max_new_tokens=300)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
+    # inputs = tokenizer(query, return_tensors="pt")
+    # outputs = model.generate(**inputs, max_new_tokens=300)
+    # response = tokenizer.decode(outputs[0]["generated_text"], skip_special_tokens=True)
+    pipe = pipeline(
+        task="text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        max_new_tokens=128,
+        return_full_text=False,
+    )
+    outputs = pipe(prompt)
+    return outputs[0]["generated_text"]
 
 # --- PDF Generation ---
 def generate_pdf(convo, topic):
@@ -159,15 +191,32 @@ if st.session_state.current_conversation:
         with st.container():
             if msg["role"] == "user":
                 if st.session_state.edit_mode.get(idx, False):
-                    new_text = st.text_input("Edit your message:", value=msg["text"], key=f"edit_input_{idx}")
+                    # Split the message into question and context
+                    question_input, context_input = msg["text"].split("<br>")
+                    # Remove the "Question:" and "Context:" parts from the beginning
+                    question_input = question_input.replace("Question: ", "")
+                    context_input = context_input.replace("Context: ", "")
+
+                    # Provide separate inputs for the question and context
+                    new_question = st.text_input("Edit your question:", value=question_input, key=f"edit_question_{idx}")
+                    new_context = st.text_area("Edit your context:", value=context_input, key=f"edit_context_{idx}")
+
+                    prompt = create_test_prompt(new_question, new_context, st.session_state.llama_tokenizer)
                     col1, col2 = st.columns([1, 1])
                     with col1:
                         if st.button("✅ Save", key=f"save_{idx}"):
-                            msg["text"] = new_text
-                            new_response = get_llama_response(new_text)
+                            # Combine question and context without the "Question:" and "Context:" labels
+                            new_combined_input = f"{new_question}<br>{new_context}"
+                            msg["text"] = new_combined_input
+                            with st.spinner("Generating response..."):
+                                try:
+                                    new_response = get_llama_response(prompt)
+                                except:
+                                    new_response = "Failed to retrieve summary."
                             if idx + 1 < len(convo) and convo[idx + 1]["role"] == "assistant":
                                 convo[idx + 1]["text"] = new_response
                             st.session_state.edit_mode[idx] = False
+                            st.session_state.chat_sessions[st.session_state.current_conversation] = convo
                             st.rerun()
                     with col2:
                         if st.button("❌ Cancel", key=f"cancel_{idx}"):
@@ -201,26 +250,42 @@ if st.session_state.current_conversation:
     if st.button("📥 Export Conversation as PDF"):
         pdf_bytes = generate_pdf(convo, st.session_state.current_conversation)
         if pdf_bytes:
-            st.download_button("Download PDF", pdf_bytes, file_name=f"{st.session_state.current_conversation}_Conversation.pdf", mime="application/pdf")
+            st.download_button("Download PDF", pdf_bytes, file_name="AskWikipedia_Conversation.pdf", mime="application/pdf")
         else:
             st.error("❌ Failed to generate PDF.")
 
     # --- User Prompt ---
-    user_input = st.chat_input("Ask Wikipedia...")
-    if user_input:
-        convo.append({"role": "user", "text": user_input})
-        st.rerun()
+    question_input = st.text_input("Enter your question:")
+    context_input = st.text_area("Enter the context:")
+
+    # Button to submit
+    if st.button("Submit"):
+        if question_input and context_input:
+            combined_input = f"Question: {question_input}<br>Context: {context_input}"
+            convo.append({"role": "user", "text": combined_input})
+
+            # Avoid rerunning unnecessarily
+            st.session_state.chat_sessions[st.session_state.current_conversation] = convo
+            st.rerun()
 
 # Display assistant response after rerun
 if st.session_state.current_conversation and len(st.session_state.chat_sessions[st.session_state.current_conversation]) % 2 == 1:
     convo = st.session_state.chat_sessions[st.session_state.current_conversation]
     last_user_msg = convo[-1]["text"]
 
+    question_input, context_input = last_user_msg.split("<br>")
+    question_input = question_input.replace("Question: ", "")
+    context_input = context_input.replace("Context: ", "")
+
+    prompt = create_test_prompt(question_input, context_input, st.session_state.llama_tokenizer)
     with st.spinner("Generating response..."):
         try:
-            assistant_reply = get_llama_response(last_user_msg)
+            assistant_reply = get_llama_response(prompt)
+            
         except Exception as e:
             assistant_reply = f"⚠️ Failed to generate response"
 
+    # Now, append the response after the spinner
     convo.append({"role": "assistant", "text": assistant_reply})
-    st.rerun()
+    st.session_state.chat_sessions[st.session_state.current_conversation] = convo
+    st.rerun()