File size: 3,928 Bytes
a7c78de d5cf00e 0b3c1f0 a7c78de 950d6d9 a7c78de 950d6d9 0b3c1f0 a7c78de 950d6d9 e944d66 9ebada2 cfcc99d 950d6d9 d5cf00e 950d6d9 9ebada2 e944d66 950d6d9 0b3c1f0 9ebada2 a7c78de e944d66 950d6d9 cfcc99d 950d6d9 cfcc99d 9ebada2 950d6d9 9ebada2 950d6d9 9ebada2 950d6d9 9ebada2 950d6d9 9ebada2 950d6d9 d5cf00e 9ebada2 950d6d9 d5cf00e 950d6d9 9ebada2 d5cf00e 9ebada2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | # app.py
# PromptPilot: Unified AI Chatbot with 5 Modes (AML-3304 Assignment)
import streamlit as st
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.markdown("AML-3304 Assignment: From Tokens to Transformers")
# Initialize session state
if 'mode' not in st.session_state:
st.session_state.mode = 'Code Generator'
# UI Buttons
cols = st.columns(5)
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("🌐 Deep Coder"):
st.session_state.mode = 'Deep Coder'
with cols[4]:
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:")
# 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
@st.cache_resource
def load_flan():
tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
return tokenizer, 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_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.write(result)
elif st.session_state.mode == 'Bayesian Q&A':
tokenizer, model = load_flan()
inputs = tokenizer(user_input, return_tensors="pt")
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 == 'Deep 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("🌐 Code Output:")
st.code(result, language="python")
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 Text:")
st.write(result)
|