Update app.py

app.py

@@ -1,39 +1,60 @@
 import streamlit as st
-from transformers import pipeline
+from langchain.chat_models import ChatHuggingFace
+from langchain.schema import SystemMessage, HumanMessage
 import torch
+from fpdf import FPDF
+import io
 
 st.set_page_config(page_title="Data Science Mentor", layout="wide")
 
-# Cache models to avoid reloading
+# Cache LangChain models per topic
 @st.cache_resource
-def load_model(topic):
-    # Select model based on topic
+def load_langchain_model(topic):
+    device = 0 if torch.cuda.is_available() else -1
     if topic == "Python":
-        return pipeline("text-generation", model="tiiuae/falcon-7b-instruct", device=0 if torch.cuda.is_available() else -1)
+        return ChatHuggingFace(repo_id="tiiuae/falcon-7b-instruct", temperature=0.6, max_new_tokens=256, task="conversational", device=device)
     elif topic == "GenAI":
-        return pipeline("text2text-generation", model="google/flan-t5-large", device=0 if torch.cuda.is_available() else -1)
+        return ChatHuggingFace(repo_id="google/flan-t5-large", temperature=0.6, max_new_tokens=256, task="conversational", device=device)
     elif topic == "Statistics":
-        return pipeline("text-generation", model="databricks/dolly-v2-3b", device=0 if torch.cuda.is_available() else -1)
+        return ChatHuggingFace(repo_id="databricks/dolly-v2-3b", temperature=0.6, max_new_tokens=256, task="conversational", device=device)
     elif topic == "SQL":
-        return pipeline("text2text-generation", model="google/flan-t5-base", device=0 if torch.cuda.is_available() else -1)
+        return ChatHuggingFace(repo_id="google/flan-t5-base", temperature=0.6, max_new_tokens=256, task="conversational", device=device)
     else:
-        # Fallback for Power BI, ML, DL
-        return pipeline("text-generation", model="tiiuae/falcon-7b-instruct", device=0 if torch.cuda.is_available() else -1)
+        # Fallback for Power BI, Machine Learning, Deep Learning
+        return ChatHuggingFace(repo_id="tiiuae/falcon-7b-instruct", temperature=0.6, max_new_tokens=256, task="conversational", device=device)
 
 def generate_answer(model, topic, level, question):
-    prompt = f"You are a {level} level mentor in {topic}. Answer the following question in detail:\n{question}"
+    system_prompt = f"You are a {level} level mentor in {topic}. Answer the user's question accordingly."
+    messages = [
+        SystemMessage(content=system_prompt),
+        HumanMessage(content=question)
+    ]
+    response = model.invoke(messages)
+    return response.content
 
-    if "text-generation" in model.task:
-        output = model(prompt, max_length=256, do_sample=True, top_k=50)
-        answer = output[0]['generated_text']
-    else:
-        output = model(prompt, max_length=256)
-        answer = output[0]['generated_text']
+def create_pdf(chat_history):
+    pdf = FPDF()
+    pdf.add_page()
+    pdf.set_auto_page_break(auto=True, margin=15)
+    pdf.set_font("Arial", size=12)
+
+    pdf.cell(0, 10, "Data Science Mentor Chat History", ln=True, align='C')
+    pdf.ln(10)
+
+    for i in range(0, len(chat_history), 2):
+        user_msg = chat_history[i][1]
+        mentor_msg = chat_history[i+1][1] if i+1 < len(chat_history) else ""
 
-    # Remove prompt from answer if echoed
-    if answer.lower().startswith(prompt.lower()):
-        answer = answer[len(prompt):].strip()
-    return answer
+        pdf.set_font("Arial", 'B', 12)
+        pdf.multi_cell(0, 10, f"You: {user_msg}")
+        pdf.set_font("Arial", '', 12)
+        pdf.multi_cell(0, 10, f"Mentor: {mentor_msg}")
+        pdf.ln(5)
+
+    pdf_buffer = io.BytesIO()
+    pdf.output(pdf_buffer)
+    pdf_buffer.seek(0)
+    return pdf_buffer
 
 # --- Streamlit UI ---
 
@@ -44,8 +65,8 @@ with st.sidebar:
     topic = st.radio("Select Topic:", ["Python", "GenAI", "Statistics", "Power BI", "SQL", "Machine Learning", "Deep Learning"])
     level = st.radio("Select Experience Level:", ["Beginner", "Intermediate", "Advanced"])
 
-# Load model for topic
-model = load_model(topic)
+# Load LangChain model for selected topic
+model = load_langchain_model(topic)
 
 if "chat_history" not in st.session_state:
     st.session_state.chat_history = []
@@ -71,5 +92,15 @@ if st.session_state.chat_history:
         st.markdown(f"**Mentor:** {mentor_msg}")
         st.markdown("---")
 
+    # PDF Download Button
+    if st.button("Download Chat as PDF"):
+        pdf_bytes = create_pdf(st.session_state.chat_history)
+        st.download_button(
+            label="Download PDF",
+            data=pdf_bytes,
+            file_name="chat_history.pdf",
+            mime="application/pdf"
+        )
+
 if st.button("Clear Chat"):
     st.session_state.chat_history = []