Update app.py

app.py

@@ -1,6 +1,5 @@
-
 # app.py
-# PromptPilot: Unified AI Chatbot with 5 Models (AML-3304 Assignment)
+# PromptPilot: Unified AI Chatbot with 5 Lightweight Modes (AML-3304 Assignment)
 
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM
@@ -14,7 +13,7 @@ st.markdown("AML-3304 Assignment: From Tokens to Transformers")
 if 'mode' not in st.session_state:
     st.session_state.mode = 'Code Generator'
 
-# UI Buttons for mode switching
+# UI Buttons
 cols = st.columns(5)
 with cols[0]:
     if st.button("🧑‍💻 Code Generator"):
@@ -26,18 +25,18 @@ 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'
+    if st.button("🌐 Lightweight Coder"):
+        st.session_state.mode = 'Lightweight Coder'
 with cols[4]:
-    if st.button("🔍 Tokenizer Introspector"):
-        st.session_state.mode = 'Tokenizer'
+    if st.button("🔁 Transformer Explorer"):
+        st.session_state.mode = 'Transformer Explorer'
 
 st.markdown(f"### 🔄 Current Mode: **{st.session_state.mode}**")
 
 # User input
 user_input = st.text_area("Enter your prompt/question:")
 
-# Cache model loaders
+# Load models
 @st.cache_resource
 def load_codegen():
     tokenizer = AutoTokenizer.from_pretrained("Salesforce/codegen-350M-mono")
@@ -50,19 +49,7 @@ def load_flan():
     model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
     return tokenizer, model
 
-@st.cache_resource
-def load_deepseek():
-    tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/deepseek-coder-1.3b-base")
-    model = AutoModelForCausalLM.from_pretrained("deepseek-ai/deepseek-coder-1.3b-base")
-    return tokenizer, model
-
-@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
-
-# Model inference
+# Inference
 if user_input and st.session_state.mode:
     with st.spinner("Generating response..."):
         if st.session_state.mode == 'Code Generator':
@@ -74,31 +61,34 @@ if user_input and st.session_state.mode:
             st.code(result, language="python")
 
         elif st.session_state.mode == 'General Q&A':
-            tokenizer, model = load_deepseek()
+            tokenizer, model = load_flan()
             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.subheader("📘 Answer:")
             st.write(result)
 
         elif st.session_state.mode == 'Bayesian Q&A':
-            tokenizer, model = load_deepseek()
+            tokenizer, model = load_flan()
             inputs = tokenizer(user_input, return_tensors="pt")
-            outputs = model.generate(**inputs, do_sample=True, temperature=0.9, top_k=40, max_new_tokens=100)
+            outputs = model.generate(**inputs, do_sample=True, top_k=40, temperature=0.9, 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()
+        elif st.session_state.mode == 'Lightweight Coder':
+            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.8)
             result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-            st.subheader("🌐 DeepSeek Code:")
+            st.subheader("🌐 Lightweight 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)
+        elif st.session_state.mode == 'Transformer Explorer':
+            tokenizer, model = load_flan()
+            prompt = f"Paraphrase this: {user_input}"
+            inputs = tokenizer(prompt, return_tensors="pt")
+            outputs = model.generate(**inputs, max_new_tokens=100)
+            result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+            st.subheader("🔁 Paraphrased by Transformer:")
+            st.write(result)