Bug fix

app.py

@@ -12,10 +12,13 @@ from langchain_core.output_parsers import StrOutputParser
 
 st.title('LLM - Retrieval Augmented Generation')
 
-models = ['tiiuae/falcon-7b-instruct', 
-          'mistralai/Mistral-7B-v0.1', 
-          'HuggingFaceH4/zephyr-7b-beta',
-          'google/gemma-7b-it']
+model_names = ['tiiuae/falcon-7b-instruct', 
+                'google/gemma-7b-it']
+
+api_urls = ['https://api-inference.huggingface.co/models/tiiuae/falcon-7b-instruct',
+            'https://api-inference.huggingface.co/models/google/gemma-7b-it']
+
+model_dict = dict(zip(model_names, api_urls))
 
 # user-input
 pdf = st.file_uploader(label='Upload PDF')
@@ -43,7 +46,7 @@ with st.sidebar:
     rag_template = st.text_area(label='Prompt template', value=default_template, height=250)
 
     st.write('# LLM parameters')
-    model = st.selectbox(label='Model', options=models, index=0)
+    model = st.selectbox(label='Model', options=model_names, index=0)
     temperature = st.slider(label='Model Temperature', min_value=0.1, max_value=float(10), value=float(1), step=0.1)
 
 # full template
@@ -119,7 +122,7 @@ def load_split_store(pdf, chunk_size, chunk_overlap):
 def format_docs(docs):
         return "\n\n".join(doc.page_content for doc in docs)
 
-@st.cache_resource
+# @st.cache_resource
 def instantiate_llm(model, temperature):
 
     # instantiate llm
@@ -133,6 +136,8 @@ def instantiate_llm(model, temperature):
         }
     )
 
+    llm.client.api_url = model_dict[model]
+
     return llm
 
 @st.cache_resource