Update app.py

app.py

@@ -1,72 +1,90 @@
 # app.py
-# Unified Multi-Model Chatbot: Programming, Bayesian Q&A, and General Q&A (AML-3304 MVP)
-
+# PromptPilot: Multi-Model AI Chatbot for AML-3304
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM
+from transformers import (AutoTokenizer, AutoModelForCausalLM,
+                          AutoModelForSeq2SeqLM, RagTokenizer, RagRetriever, RagSequenceForGeneration)
 import torch
 
-st.set_page_config(page_title="PromptPilot", layout="centered")
-st.title("🤖 PromptPilot")
+st.set_page_config(page_title="PromptPilot: Unified AI Chatbot", layout="centered")
+st.title("🤖 PromptPilot: Unified AI Chatbot")
 
-st.markdown("DeepSeek-R1, FLAN-T5, CodeGen, and Bayesian-style demos unified in one app ✅")
+st.markdown("""
+**MVP Chatbot Demo for AML-3304**
+Supports:
+- Basic Code Generation
+- General Q&A
+- Bayesian-style Q&A
+- Creative Writing
+- DeepSeek‑R1 RAG (Retrieval-Augmented Q&A)
+""")
 
-# Track mode with session state
 if 'mode' not in st.session_state:
     st.session_state.mode = 'Code Generator'
 
-# Layout buttons in a row
-col1, col2, col3 = st.columns(3)
-with col1:
-    if st.button("🧑‍💻 Code Generator"):
-        st.session_state.mode = 'Code Generator'
-with col2:
-    if st.button("📘 General Q&A"):
-        st.session_state.mode = 'General Q&A'
-with col3:
-    if st.button("📊 Bayesian Q&A"):
-        st.session_state.mode = 'Bayesian Q&A'
+cols = st.columns(5)
+labels = [
+    ("🧑‍💻 Basic Code", 'Code Generator'),
+    ("📘 General Q&A", 'General Q&A'),
+    ("📊 Bayesian Q&A", 'Bayesian Q&A'),
+    ("✍️ Creative Writing", 'Creative Writing'),
+    ("🔍 RAG Q&A", 'RAG Q&A'),
+]
+for col, (emoji, mode_name) in zip(cols, labels):
+    if col.button(emoji):
+        st.session_state.mode = mode_name
 
 st.markdown(f"### 🔄 Current Mode: **{st.session_state.mode}**")
+user_input = st.text_input("Enter your prompt/question:")
 
-# Text input
-user_input = st.text_input("Enter your question or prompt below:")
-
-# Load models
 @st.cache_resource
 def load_codegen():
-    tokenizer = AutoTokenizer.from_pretrained("Salesforce/codegen-350M-mono")
-    model = AutoModelForCausalLM.from_pretrained("Salesforce/codegen-350M-mono")
-    return tokenizer, model
+    return (AutoTokenizer.from_pretrained("Salesforce/codegen-350M-mono"),
+            AutoModelForCausalLM.from_pretrained("Salesforce/codegen-350M-mono"))
+
+@st.cache_resource
+def load_flan():
+    return (AutoTokenizer.from_pretrained("google/flan-t5-base"),
+            AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base"))
 
 @st.cache_resource
-def load_flan_t5():
-    tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
-    model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
+def load_falcon():
+    return (AutoTokenizer.from_pretrained("tiiuae/falcon-rw-1b"),
+            AutoModelForCausalLM.from_pretrained("tiiuae/falcon-rw-1b"))
+
+@st.cache_resource
+def load_rag():
+    tokenizer = RagTokenizer.from_pretrained("deepseek-ai/deepseek-r1-rag")
+    retriever = RagRetriever.from_pretrained("deepseek-ai/deepseek-r1-rag", index_name="custom")
+    model = RagSequenceForGeneration.from_pretrained("deepseek-ai/deepseek-r1-rag", retriever=retriever)
     return tokenizer, model
 
-# Generate response
-if user_input and st.session_state.mode:
+if user_input:
     with st.spinner("Generating response..."):
-        if st.session_state.mode == 'Code Generator':
-            tokenizer, model = load_codegen()
-            inputs = tokenizer(user_input, return_tensors="pt")
-            outputs = model.generate(inputs["input_ids"], max_new_tokens=128, do_sample=True, temperature=0.7, pad_token_id=tokenizer.eos_token_id)
-            result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-            st.subheader("💻 Generated Code:")
-            st.code(result, language="python")
+        mode = st.session_state.mode
+
+        if mode == 'Code Generator':
+            tok, mod = load_codegen()
+            inp = tok(user_input, return_tensors="pt")
+            out = mod.generate(inp.input_ids, max_new_tokens=128, temperature=0.7, do_sample=True)
+            st.code(tok.decode(out[0], skip_special_tokens=True), language="python")
+
+        elif mode in ('General Q&A', 'Bayesian Q&A'):
+            tok, mod = load_flan()
+            inp = tok(user_input, return_tensors="pt")
+            if mode == 'General Q&A':
+                out = mod.generate(inp.input_ids, max_new_tokens=100)
+            else:
+                out = mod.generate(inp.input_ids, max_new_tokens=100, do_sample=True, temperature=0.9, top_k=50)
+            st.write(tok.decode(out[0], skip_special_tokens=True))
 
-        elif st.session_state.mode == 'General Q&A':
-            tokenizer, model = load_flan_t5()
-            inputs = tokenizer(user_input, return_tensors="pt")
-            outputs = model.generate(inputs["input_ids"], max_new_tokens=100)
-            result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-            st.subheader("📘 Answer:")
-            st.write(result)
+        elif mode == 'Creative Writing':
+            tok, mod = load_falcon()
+            inp = tok(user_input, return_tensors="pt")
+            out = mod.generate(inp.input_ids, max_new_tokens=150, do_sample=True, top_p=0.95, temperature=1.0)
+            st.write(tok.decode(out[0], skip_special_tokens=True))
 
-        elif st.session_state.mode == 'Bayesian Q&A':
-            tokenizer, model = load_flan_t5()
-            inputs = tokenizer(user_input, return_tensors="pt")
-            outputs = model.generate(inputs["input_ids"], do_sample=True, temperature=0.9, top_k=50, max_new_tokens=100)
-            result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-            st.subheader("📊 Bayesian-style Answer (Sampled):")
-            st.write(result)
+        elif mode == 'RAG Q&A':
+            tok, mod = load_rag()
+            inp = tok(user_input, return_tensors="pt")
+            out = mod.generate(input_ids=inp["input_ids"], context_input_ids=inp["context_input_ids"], max_new_tokens=100)
+            st.write(tok.batch_decode(out, skip_special_tokens=True)[0])