Spaces:
Running
Running
File size: 7,411 Bytes
62098f3 c2d1e89 62098f3 409d83c f0b46d5 62098f3 c2d1e89 62098f3 c4c458c 62098f3 409d83c 62098f3 d371388 62098f3 c2d1e89 62098f3 c4c458c f0b46d5 62098f3 c4c458c f70af54 732cf88 c2d1e89 732cf88 f0b46d5 732cf88 f0b46d5 732cf88 c2d1e89 c4c458c 62098f3 f70af54 e559684 |
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 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 |
import streamlit as st
from transformers import T5ForConditionalGeneration, T5Tokenizer, pipeline, AutoTokenizer, AutoModelForCausalLM
import torch
# Streamlit app setup
st.set_page_config(page_title="Chat", layout="wide")
# Sidebar: Model controls
st.sidebar.title("Model Controls")
model_options = {
"1": "karthikeyan-r/slm-custom-model_6k",
"2": "karthikeyan-r/calculation_model"
}
model_choice = st.sidebar.selectbox("Select Model", options=list(model_options.values()))
load_model_button = st.sidebar.button("Load Model")
clear_conversation_button = st.sidebar.button("Clear Conversation")
clear_model_button = st.sidebar.button("Clear Model")
# Main UI
st.title("Chat Conversation UI")
# Session states
if "model" not in st.session_state:
st.session_state["model"] = None
if "tokenizer" not in st.session_state:
st.session_state["tokenizer"] = None
if "qa_pipeline" not in st.session_state:
st.session_state["qa_pipeline"] = None
if "conversation" not in st.session_state:
st.session_state["conversation"] = []
if "user_input" not in st.session_state:
st.session_state["user_input"] = ""
# Load Model
if load_model_button:
with st.spinner("Loading model..."):
try:
# Load the selected model
if model_choice == model_options["1"]:
# Load the T5 model for general QA (slm-custom-model_6k)
device = 0 if torch.cuda.is_available() else -1
st.session_state["model"] = T5ForConditionalGeneration.from_pretrained(model_choice, cache_dir="./model_cache")
st.session_state["tokenizer"] = T5Tokenizer.from_pretrained(model_choice, cache_dir="./model_cache")
st.session_state["qa_pipeline"] = pipeline(
"text2text-generation",
model=st.session_state["model"],
tokenizer=st.session_state["tokenizer"],
device=device
)
elif model_choice == model_options["2"]:
# Load the calculation model (calculation_model)
tokenizer = AutoTokenizer.from_pretrained(model_choice, cache_dir="./model_cache")
model = AutoModelForCausalLM.from_pretrained(model_choice, cache_dir="./model_cache")
# Add special tokens if not present
if tokenizer.pad_token is None:
tokenizer.add_special_tokens({'pad_token': '[PAD]'})
model.resize_token_embeddings(len(tokenizer))
if tokenizer.eos_token is None:
tokenizer.add_special_tokens({'eos_token': '[EOS]'})
model.resize_token_embeddings(len(tokenizer))
# Update configuration
model.config.pad_token_id = tokenizer.pad_token_id
model.config.eos_token_id = tokenizer.eos_token_id
st.session_state["model"] = model
st.session_state["tokenizer"] = tokenizer
st.session_state["qa_pipeline"] = None # Calculation model doesn't use text2text pipeline
st.success("Model loaded successfully and ready!")
except Exception as e:
st.error(f"Error loading model: {e}")
# Clear Model
if clear_model_button:
st.session_state["model"] = None
st.session_state["tokenizer"] = None
st.session_state["qa_pipeline"] = None
st.success("Model cleared.")
# Chat Conversation Display
def display_conversation():
"""Display the chat conversation dynamically."""
st.subheader("Conversation")
for idx, (speaker, message) in enumerate(st.session_state["conversation"]):
if speaker == "You":
st.markdown(f"**You:** {message}")
else:
st.markdown(f"**Model:** {message}")
display_conversation()
# Input Area
if st.session_state["qa_pipeline"]:
user_input = st.text_input(
"Enter your query:",
value=st.session_state["user_input"], # Use session state for persistence
key="chat_input",
)
if st.button("Send", key="send_button"):
if user_input:
with st.spinner("Generating response..."):
try:
# Generate the model response for general QA (T5 model)
response = st.session_state["qa_pipeline"](f"Q: {user_input}", max_length=400)
generated_text = response[0]["generated_text"]
# Update the conversation
st.session_state["conversation"].append(("You", user_input))
st.session_state["conversation"].append(("Model", generated_text))
# Clear the input field after submission
st.session_state["user_input"] = ""
# Rerender the conversation immediately
display_conversation()
except Exception as e:
st.error(f"Error generating response: {e}")
else:
# Handle user input for the calculation model (calculation_model)
if st.session_state["model"] and model_choice == model_options["2"]:
user_input = st.text_input(
"Enter your query for calculation:",
value=st.session_state["user_input"],
key="calculation_input",
)
if st.button("Send Calculation", key="send_calculation_button"):
if user_input:
with st.spinner("Generating response..."):
try:
# Generate the model response for the calculation model
inputs = st.session_state["tokenizer"](f"Input: {user_input}\nOutput:", return_tensors="pt", padding=True, truncation=True)
input_ids = inputs.input_ids
attention_mask = inputs.attention_mask
output = st.session_state["model"].generate(
input_ids=input_ids,
attention_mask=attention_mask,
max_length=50,
pad_token_id=st.session_state["tokenizer"].pad_token_id,
eos_token_id=st.session_state["tokenizer"].eos_token_id,
do_sample=False
)
decoded_output = st.session_state["tokenizer"].decode(output[0], skip_special_tokens=True)
if "Output:" in decoded_output:
answer = decoded_output.split("Output:")[-1].strip()
else:
answer = decoded_output.strip()
# Update the conversation
st.session_state["conversation"].append(("You", user_input))
st.session_state["conversation"].append(("Model", answer))
# Clear the input field after submission
st.session_state["user_input"] = ""
# Rerender the conversation immediately
display_conversation()
except Exception as e:
st.error(f"Error generating response: {e}")
# Clear Conversation
if clear_conversation_button:
st.session_state["conversation"] = []
st.session_state["user_input"] = "" # Clear input field
st.success("Conversation cleared.")
|