Update app.py

app.py

@@ -1,117 +1,103 @@
 # app.py
-# PromptPilot: Unified Multi-Model Chatbot (AML-3304 Final Submission)
+# PromptPilot: Unified AI Chatbot with 5 Models (AML-3304 Assignment)
 
 import streamlit as st
-from transformers import (
-    AutoTokenizer,
-    AutoModelForCausalLM,
-    AutoModelForSeq2SeqLM
-)
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM
 import torch
 
-st.set_page_config(page_title="PromptPilot", layout="centered")
+st.set_page_config(page_title="PromptPilot - Unified AI Chatbot", layout="centered")
 st.title("🤖 PromptPilot - Unified AI Chatbot")
-
 st.markdown("AML-3304 Assignment: From Tokens to Transformers")
 
-# Track current mode
-if "mode" not in st.session_state:
-    st.session_state.mode = "Code Generator"
+# Initialize session state
+if 'mode' not in st.session_state:
+    st.session_state.mode = 'Code Generator'
 
-# Mode selection buttons
+# UI Buttons for mode switching
 cols = st.columns(5)
-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 i, (label, mode) in enumerate(mode_labels):
-    if cols[i].button(label):
-        st.session_state.mode = mode
+with cols[0]:
+    if st.button("🧑‍💻 Code Generator"):
+        st.session_state.mode = 'Code Generator'
+with cols[1]:
+    if st.button("📘 General Q&A"):
+        st.session_state.mode = 'General Q&A'
+with cols[2]:
+    if st.button("📊 Bayesian Q&A"):
+        st.session_state.mode = 'Bayesian Q&A'
+with cols[3]:
+    if st.button("🌐 DeepSeek Coder"):
+        st.session_state.mode = 'DeepSeek Coder'
+with cols[4]:
+    if st.button("🔍 Tokenizer Introspector"):
+        st.session_state.mode = 'Tokenizer'
 
 st.markdown(f"### 🔄 Current Mode: **{st.session_state.mode}**")
 
-# Input field
-user_input = st.text_area("Enter your prompt/question:", height=100)
+# User input
+user_input = st.text_area("Enter your prompt/question:")
 
-# Cached model loaders
+# Cache model loaders
 @st.cache_resource
 def load_codegen():
-    tok = AutoTokenizer.from_pretrained("Salesforce/codegen-350M-mono")
-    mod = AutoModelForCausalLM.from_pretrained("Salesforce/codegen-350M-mono")
-    return tok, mod
+    tokenizer = AutoTokenizer.from_pretrained("Salesforce/codegen-350M-mono")
+    model = AutoModelForCausalLM.from_pretrained("Salesforce/codegen-350M-mono")
+    return tokenizer, model
 
 @st.cache_resource
 def load_flan():
-    tok = AutoTokenizer.from_pretrained("google/flan-t5-base")
-    mod = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
-    return tok, mod
+    tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
+    model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
+    return tokenizer, model
 
 @st.cache_resource
 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
+    tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/deepseek-llm-7b-instruct")
+    model = AutoModelForCausalLM.from_pretrained("deepseek-ai/deepseek-llm-7b-instruct")
+    return tokenizer, model
 
-# Run selected model logic
-if user_input:
-    with st.spinner("Generating response..."):
-        if st.session_state.mode == "Code Generator":
-            tok, mod = load_codegen()
-            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")
+@st.cache_resource
+def load_deepseek_coder():
+    tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/deepseek-coder-6.7b-instruct")
+    model = AutoModelForCausalLM.from_pretrained("deepseek-ai/deepseek-coder-6.7b-instruct")
+    return tokenizer, model
 
-        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)
+# Model inference
+if user_input and st.session_state.mode:
+    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")
+
+        elif st.session_state.mode == 'General Q&A':
+            tokenizer, model = load_deepseek()
+            inputs = tokenizer(user_input, return_tensors="pt")
+            outputs = model.generate(**inputs, max_new_tokens=150)
+            result = tokenizer.decode(outputs[0], skip_special_tokens=True)
             st.subheader("📘 Answer")
-            st.write(res)
-
-        elif st.session_state.mode == "Bayesian Q&A":
-            tok, mod = load_flan()
-            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})
+            st.write(result)
+
+        elif st.session_state.mode == 'Bayesian Q&A':
+            tokenizer, model = load_deepseek()
+            inputs = tokenizer(user_input, return_tensors="pt")
+            outputs = model.generate(**inputs, do_sample=True, temperature=0.9, top_k=40, max_new_tokens=100)
+            result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+            st.subheader("📊 Bayesian-style Answer:")
+            st.write(result)
+
+        elif st.session_state.mode == 'DeepSeek Coder':
+            tokenizer, model = load_deepseek_coder()
+            inputs = tokenizer(user_input, return_tensors="pt")
+            outputs = model.generate(inputs["input_ids"], max_new_tokens=128, do_sample=True, temperature=0.8)
+            result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+            st.subheader("🌐 DeepSeek Code:")
+            st.code(result, language="python")
+
+        elif st.session_state.mode == 'Tokenizer':
+            tokenizer, _ = load_flan()
+            tokens = tokenizer.tokenize(user_input)
+            st.subheader("🔬 Tokenized Output:")
+            st.write(tokens)