Update app.py

app.py

@@ -1,25 +1,36 @@
 import os
+import subprocess
 import streamlit as st
 from dotenv import load_dotenv
-import wikipedia
 from xhtml2pdf import pisa
 import io
-
-# Load environment variables
-load_dotenv()
+from transformers import AutoTokenizer, AutoModelForCausalLM  # 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")
+    return model, tokenizer
+
+# --- 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
+    st.session_state.llama_tokenizer = tokenizer
 
 st.set_page_config(page_title="Ask Wikipedia", page_icon="📘", layout="wide")
 
-# --- Wikipedia Summary ---
-def get_wikipedia_summary(query):
-    try:
-        return wikipedia.summary(query, sentences=2)
-    except wikipedia.exceptions.DisambiguationError as e:
-        return f"Your query is ambiguous, here are some options: {e.options}"
-    except wikipedia.exceptions.HTTPTimeoutError:
-        return "Request timed out. Please try again later."
-    except Exception as e:
-        return f"An error occurred: {e}"
+def get_llama_response(query):
+    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
 
 # --- PDF Generation ---
 def generate_pdf(convo, topic):
@@ -153,10 +164,7 @@ if st.session_state.current_conversation:
                     with col1:
                         if st.button("✅ Save", key=f"save_{idx}"):
                             msg["text"] = new_text
-                            try:
-                                new_response = get_wikipedia_summary(new_text)
-                            except:
-                                new_response = "Failed to retrieve summary."
+                            new_response = get_llama_response(new_text)
                             if idx + 1 < len(convo) and convo[idx + 1]["role"] == "assistant":
                                 convo[idx + 1]["text"] = new_response
                             st.session_state.edit_mode[idx] = False
@@ -191,9 +199,9 @@ if st.session_state.current_conversation:
 
     # --- Export PDF ---
     if st.button("📥 Export Conversation as PDF"):
-        pdf_bytes = generate_pdf(convo,st.session_state.current_conversation)
+        pdf_bytes = generate_pdf(convo, st.session_state.current_conversation)
         if pdf_bytes:
-            st.download_button("Download PDF", pdf_bytes, file_name="AskWikipedia_Conversation.pdf", mime="application/pdf")
+            st.download_button("Download PDF", pdf_bytes, file_name=f"{st.session_state.current_conversation}_Conversation.pdf", mime="application/pdf")
         else:
             st.error("❌ Failed to generate PDF.")
 
@@ -201,9 +209,18 @@ if st.session_state.current_conversation:
     user_input = st.chat_input("Ask Wikipedia...")
     if user_input:
         convo.append({"role": "user", "text": user_input})
+        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"]
+
+    with st.spinner("Generating response..."):
         try:
-            reply = get_wikipedia_summary(user_input)
-        except:
-            reply = "Could not fetch response."
-        convo.append({"role": "assistant", "text": reply})
-        st.rerun()
+            assistant_reply = get_llama_response(last_user_msg)
+        except Exception as e:
+            assistant_reply = f"⚠️ Failed to generate response"
+
+    convo.append({"role": "assistant", "text": assistant_reply})
+    st.rerun()