Update src/streamlit_app.py

src/streamlit_app.py

@@ -2,26 +2,31 @@ import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 # Load your model and tokenizer
-model_name = "CJHauser/PrisimAI-chat"  # Your Hugging Face model
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+model_name = "CJHauser/PrisimAI-chat"
+tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
 # Streamlit App
-st.title("PrisimAI Chatbot")
+st.set_page_config(page_title="PrisimAI Chatbot")
+st.title("🤖 PrisimAI Chatbot")
 
-# Input box for user prompt
-user_input = st.text_input("Ask a question:")
+# User prompt
+user_input = st.text_input("Ask something:", placeholder="e.g. What is AI?")
 
-# When the user enters a prompt
 if user_input:
-    # Encode the input
-    inputs = tokenizer(user_input, return_tensors="pt")
-    
-    # Generate the response
-    outputs = model.generate(inputs['input_ids'], max_length=100)
-    
-    # Decode the output
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    # Show the response
-    st.write(f"Response: {response}")
+    with st.spinner("Thinking..."):
+        inputs = tokenizer(user_input, return_tensors="pt")
+        outputs = model.generate(
+            inputs["input_ids"],
+            max_length=150,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9,
+            pad_token_id=tokenizer.eos_token_id
+        )
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    # Remove the original prompt from the response if repeated
+    response = response.replace(user_input, "").strip()
+
+    st.markdown(f"**Response:** {response}")