diff --git a/src/components/ChatDemo/search.tsx b/src/components/ChatDemo/search.tsx
index e12a4d0..95851a6 100644
--- a/src/components/ChatDemo/search.tsx
+++ b/src/components/ChatDemo/search.tsx
@@ -57,11 +57,11 @@ export const Search: FC<SearchProps> = ({ pipeline, setQuery }) => {
           }
           autoFocus
           placeholder="Question everything..."
-          className="w-full px-4 py-2 bg-zinc-700 text-zinc-200 rounded-l-lg focus:outline-none focus:ring-2 focus:ring-blue-500"
+          className="w-full px-4 py-2 h-10 bg-zinc-700 text-zinc-200 rounded-l-lg focus:outline-none focus:ring-2 focus:ring-blue-500"
         />
         <button
           type="submit"
-          className="px-4 py-2 bg-blue-500 text-white rounded-r-lg hover:bg-blue-600 focus:outline-none focus:ring-2 focus:ring-blue-500"
+          className="px-4 py-2 h-10 bg-blue-500 text-white rounded-r-lg hover:bg-blue-600 focus:outline-none focus:ring-2 focus:ring-blue-500"
         >
           <ArrowRight size={20} />
         </button>
diff --git a/src/components/ui/ModelSelector.tsx b/src/components/ui/ModelSelector.tsx
index b155cbe..f30d131 100644
--- a/src/components/ui/ModelSelector.tsx
+++ b/src/components/ui/ModelSelector.tsx
@@ -14,6 +14,7 @@ import {
   SelectTrigger,
   SelectValue,
 } from '@/components/ui/select';
+import { useUserContext } from '@/context/UserContext';
 
 interface ModelSelectorProps {
   id?: string;
@@ -21,11 +22,8 @@ interface ModelSelectorProps {
   setSelectedModel: (model: string) => void;
 }
 
-const ModelSelector: React.FC<ModelSelectorProps> = ({
-  id,
-  selectedModel,
-  setSelectedModel,
-}) => {
+const ModelSelector: React.FC<{ id?: string }> = ({ id }) => {
+  const { selectedModel, setSelectedModel } = useUserContext();
   const [isDialogOpen, setIsDialogOpen] = useState(false);
   const [customModelValue, setCustomModelValue] = useState('');
 
@@ -56,10 +54,10 @@ const ModelSelector: React.FC<ModelSelectorProps> = ({
             <SelectValue>{selectedModel || 'Select a model'}</SelectValue>
           </SelectTrigger>
           <SelectContent>
+            <SelectItem value="gpt-4o">GPT-4o</SelectItem>
             <SelectItem value="gpt-4-turbo">GPT-4 Turbo</SelectItem>
-            <SelectItem value="gpt-4">GPT-4</SelectItem>
             <SelectItem value="gpt-3.5-turbo">GPT-3.5 Turbo</SelectItem>
-            <SelectItem value="llama-2">Llama-2</SelectItem>
+            <SelectItem value="ollama/llama2">Llama 2</SelectItem>
             <SelectItem value="customModel">Add another model</SelectItem>
           </SelectContent>
         </Select>
diff --git a/src/context/UserContext.tsx b/src/context/UserContext.tsx
index 1dd847d..868b437 100644
--- a/src/context/UserContext.tsx
+++ b/src/context/UserContext.tsx
@@ -15,6 +15,8 @@ interface UserContextProps {
     name: string,
     url: string
   ) => { nameUnique: boolean; urlUnique: boolean };
+  selectedModel: string;
+  setSelectedModel: (model: string) => void;
 }
 
 const UserContext = createContext<UserContextProps>({
@@ -22,6 +24,8 @@ const UserContext = createContext<UserContextProps>({
   addWatchedPipeline: () => {},
   removeWatchedPipeline: () => {},
   isPipelineUnique: () => ({ nameUnique: true, urlUnique: true }),
+  selectedModel: 'gpt-4o',
+  setSelectedModel: () => {},
 });
 
 export const useUserContext = () => useContext(UserContext);
@@ -39,6 +43,13 @@ export const UserProvider: React.FC<{ children: React.ReactNode }> = ({
     },
   });
 
+  const [selectedModel, setSelectedModel] = useState(() => {
+    if (typeof window !== 'undefined') {
+      return localStorage.getItem('selectedModel') || 'gpt-4-turbo';
+    }
+    return 'gpt-4-turbo';
+  });
+
   // Load data from local storage on initial render
   useEffect(() => {
     const storedPipelines = localStorage.getItem('watchedPipelines');
@@ -59,6 +70,11 @@ export const UserProvider: React.FC<{ children: React.ReactNode }> = ({
     localStorage.setItem('watchedPipelines', JSON.stringify(watchedPipelines));
   }, [watchedPipelines]);
 
+  // Save selectedModel to local storage whenever it changes
+  useEffect(() => {
+    localStorage.setItem('selectedModel', selectedModel);
+  }, [selectedModel]);
+
   const addWatchedPipeline = (pipelineId: string, pipeline: Pipeline) => {
     setWatchedPipelines((prevPipelines) => ({
       ...prevPipelines,
@@ -91,6 +107,8 @@ export const UserProvider: React.FC<{ children: React.ReactNode }> = ({
         addWatchedPipeline,
         removeWatchedPipeline,
         isPipelineUnique,
+        selectedModel,
+        setSelectedModel,
       }}
     >
       {children}
diff --git a/src/pages/pipeline/[pipelineId]/playground.tsx b/src/pages/pipeline/[pipelineId]/playground.tsx
index 5f02714..714dd7f 100644
--- a/src/pages/pipeline/[pipelineId]/playground.tsx
+++ b/src/pages/pipeline/[pipelineId]/playground.tsx
@@ -1,7 +1,5 @@
 'use client';
-import { Loader } from 'lucide-react';
 import { useSearchParams } from 'next/navigation';
-import { useRouter } from 'next/router';
 import { r2rClient } from 'r2r-js';
 import React, { useState, useEffect, useRef } from 'react';
 
@@ -9,7 +7,6 @@ import { Result } from '@/components/ChatDemo/result';
 import { Search } from '@/components/ChatDemo/search';
 import SingleSwitch from '@/components/ChatDemo/SingleSwitch';
 import useSwitchManager from '@/components/ChatDemo/SwitchManager';
-import { Title } from '@/components/ChatDemo/title';
 import Layout from '@/components/Layout';
 import { Input } from '@/components/ui/input';
 import { Label } from '@/components/ui/label';
@@ -17,10 +14,10 @@ import ModelSelector from '@/components/ui/ModelSelector';
 import { Slider } from '@/components/ui/slider';
 import { useToast } from '@/components/ui/use-toast';
 import { usePipelineInfo } from '@/context/PipelineInfo';
+import { useUserContext } from '@/context/UserContext';
 
 const Index: React.FC = () => {
   const searchParams = useSearchParams();
-  const router = useRouter();
   const { toast } = useToast();
   const [query, setQuery] = useState('');
   const [hasAttemptedFetch, setHasAttemptedFetch] = useState(false);
@@ -31,8 +28,7 @@ const Index: React.FC = () => {
     }
   }, [searchParams]);
 
-  const [sidebarWidth, setSidebarWidth] = useState(300);
-  const [selectedModel, setSelectedModel] = useState('gpt-4-turbo');
+  const { selectedModel } = useUserContext();
   const [uploadedDocuments, setUploadedDocuments] = useState([]);
 
   const { pipeline, isLoading: isPipelineLoading } = usePipelineInfo();
@@ -158,13 +154,8 @@ const Index: React.FC = () => {
               <div className="flex flex-col gap-2">
                 <div className="flex flex-col gap-2">
                   <Label htmlFor="selectedModel">Selected Model</Label>
-                  <Input
-                    id="selectedModel"
-                    type="string"
-                    value={selectedModel}
-                    // className="w-24"
-                    onChange={(e) => setSelectedModel(e.target.value)}
-                  />
+                  {/* Model Selector */}
+                  <ModelSelector id={selectedModel} />
                 </div>
                 <div className="flex flex-col gap-2">
                   <Label htmlFor="top_k">Top K</Label>
@@ -230,11 +221,6 @@ const Index: React.FC = () => {
               disabled={true}
               className="w-full bg-zinc-700 text-zinc-300 p-2 rounded"
             />
-            {/* Model Selector */}
-            {/* <ModelSelector
-              selectedModel={selectedModel}
-              setSelectedModel={setSelectedModel}
-            /> */}
           </div>
         </div>