Update app.py

app.py

@@ -1,90 +1,117 @@
 # app.py
-# PromptPilot: Multi-Model AI Chatbot for AML-3304
+# PromptPilot: Unified Multi-Model Chatbot (AML-3304 Final Submission)
+
 import streamlit as st
-from transformers import (AutoTokenizer, AutoModelForCausalLM,
-                          AutoModelForSeq2SeqLM, RagTokenizer, RagRetriever, RagSequenceForGeneration)
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    AutoModelForSeq2SeqLM
+)
 import torch
 
-st.set_page_config(page_title="PromptPilot: Unified AI Chatbot", layout="centered")
-st.title("🤖 PromptPilot: Unified AI Chatbot")
+st.set_page_config(page_title="PromptPilot", layout="centered")
+st.title("🤖 PromptPilot - Unified AI Chatbot")
 
-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)
-""")
+st.markdown("AML-3304 Assignment: From Tokens to Transformers")
 
-if 'mode' not in st.session_state:
-    st.session_state.mode = 'Code Generator'
+# Track current mode
+if "mode" not in st.session_state:
+    st.session_state.mode = "Code Generator"
 
+# Mode selection buttons
 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'),
+mode_labels = [
+    ("👨‍💻 Code Generator", "Code Generator"),
+    ("📘 General Q&A", "General Q&A"),
+    ("📊 Bayesian Q&A", "Bayesian Q&A"),
+    ("🌐 DeepSeek Coder", "DeepSeek"),
+    ("🔍 Tokenizer Introspector", "Tokenizer")
 ]
-for col, (emoji, mode_name) in zip(cols, labels):
-    if col.button(emoji):
-        st.session_state.mode = mode_name
+
+for i, (label, mode) in enumerate(mode_labels):
+    if cols[i].button(label):
+        st.session_state.mode = mode
 
 st.markdown(f"### 🔄 Current Mode: **{st.session_state.mode}**")
-user_input = st.text_input("Enter your prompt/question:")
 
+# Input field
+user_input = st.text_area("Enter your prompt/question:", height=100)
+
+# Cached model loaders
 @st.cache_resource
 def load_codegen():
-    return (AutoTokenizer.from_pretrained("Salesforce/codegen-350M-mono"),
-            AutoModelForCausalLM.from_pretrained("Salesforce/codegen-350M-mono"))
+    tok = AutoTokenizer.from_pretrained("Salesforce/codegen-350M-mono")
+    mod = AutoModelForCausalLM.from_pretrained("Salesforce/codegen-350M-mono")
+    return tok, mod
 
 @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_falcon():
-    return (AutoTokenizer.from_pretrained("tiiuae/falcon-rw-1b"),
-            AutoModelForCausalLM.from_pretrained("tiiuae/falcon-rw-1b"))
+    tok = AutoTokenizer.from_pretrained("google/flan-t5-base")
+    mod = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
+    return tok, mod
 
 @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
+def load_deepseek():
+    tok = AutoTokenizer.from_pretrained("deepseek-ai/deepseek-coder-1.3b-instruct")
+    mod = AutoModelForCausalLM.from_pretrained("deepseek-ai/deepseek-coder-1.3b-instruct")
+    return tok, mod
 
+# Run selected model logic
 if user_input:
     with st.spinner("Generating response..."):
-        mode = st.session_state.mode
-
-        if mode == 'Code Generator':
+        if st.session_state.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")
+            inputs = tok(user_input, return_tensors="pt")
+            outputs = mod.generate(
+                inputs.input_ids,
+                max_new_tokens=128,
+                do_sample=True,
+                temperature=0.7,
+                pad_token_id=tok.eos_token_id
+            )
+            res = tok.decode(outputs[0], skip_special_tokens=True)
+            st.subheader("💻 Code Output")
+            st.code(res, language="python")
+
+        elif st.session_state.mode == "General Q&A":
+            tok, mod = load_flan()
+            inputs = tok(user_input, return_tensors="pt")
+            outputs = mod.generate(inputs.input_ids, max_new_tokens=100)
+            res = tok.decode(outputs[0], skip_special_tokens=True)
+            st.subheader("📘 Answer")
+            st.write(res)
 
-        elif mode in ('General Q&A', 'Bayesian Q&A'):
+        elif st.session_state.mode == "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 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 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])
+            inputs = tok(user_input, return_tensors="pt")
+            outputs = mod.generate(
+                inputs.input_ids,
+                do_sample=True,
+                top_k=50,
+                temperature=0.9,
+                max_new_tokens=100
+            )
+            res = tok.decode(outputs[0], skip_special_tokens=True)
+            st.subheader("📊 Bayesian-style Sampled Answer")
+            st.write(res)
+
+        elif st.session_state.mode == "DeepSeek":
+            tok, mod = load_deepseek()
+            inputs = tok(user_input, return_tensors="pt")
+            outputs = mod.generate(
+                inputs.input_ids,
+                max_new_tokens=128,
+                temperature=0.7,
+                do_sample=True,
+                pad_token_id=tok.eos_token_id
+            )
+            res = tok.decode(outputs[0], skip_special_tokens=True)
+            st.subheader("🌐 DeepSeek Coder Output")
+            st.code(res, language="python")
+
+        elif st.session_state.mode == "Tokenizer":
+            tok, _ = load_flan()
+            tokens = tok.tokenize(user_input)
+            ids = tok.convert_tokens_to_ids(tokens)
+            st.subheader("🔍 Tokenization Breakdown")
+            st.write({"Tokens": tokens, "Token IDs": ids})