Initial Draft

app.py

@@ -3,7 +3,6 @@ import tempfile
 import os
 import re
 import torch
-from threading import Thread
 
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer
 from langchain_community.document_loaders import PyPDFLoader
@@ -52,23 +51,22 @@ def fn_generate_QnA_response(mv_selected_model, mv_user_question, lv_vector_stor
     print("Step4: Generating LLM response")
     mv_processing_message.text("Step4: Generating LLM response")
 
-    lv_tokenizer = AutoTokenizer.from_pretrained(mv_selected_model, trust_remote_code=True)
-    lv_model = AutoModelForCausalLM.from_pretrained(
+    lv_tokenizer = AutoTokenizer.from_pretrained(
                                                         mv_selected_model, 
-                                                        torch_dtype="auto", 
+                                                        model_max_length=2048,
+                                                        trust_remote_code=True
+                                                )
+    lv_model = AutoModelForCausalLM.from_pretrained(
+                                                        mv_selected_model,  
                                                         device_map="cpu", 
                                                         trust_remote_code=True
                                                    )
-    # lv_streamer = TextIteratorStreamer(
-    #                                         tokenizer=lv_tokenizer, 
-    #                                         skip_prompt=True, 
-    #                                         skip_special_tokens=True, 
-    #                                         timeout=300.0
-    #                                   )
     lv_ms_phi2_pipeline = pipeline(
                                     "text-generation", tokenizer=lv_tokenizer, model=lv_model,
-                                    device_map="cpu", max_new_tokens=512, return_full_text=True
+                                    pad_token_id=lv_tokenizer.eos_token_id, eos_token_id=lv_tokenizer.eos_token_id,
+                                    device_map="cpu", max_new_tokens=2048, return_full_text=True
                                   )
+                                  
     lv_hf_phi2_pipeline = HuggingFacePipeline(pipeline=lv_ms_phi2_pipeline)
     lv_chain = ConversationalRetrievalChain.from_llm(lv_hf_phi2_pipeline, lv_vector_store.as_retriever(), return_source_documents=True)
     lv_response = lv_chain({"question": mv_user_question, 'chat_history': lv_chat_history})
@@ -116,9 +114,6 @@ def main():
     st.text("")
     st.text("")
     st.text("")
-    st.text("")
-    st.text("")
-    st.text("")
 
     mv_vector_storage_dir = "/workspace/knowledge-base/01-ML/01-dev/adhoc/Talk2PDF/vector_store"
     
@@ -164,7 +159,6 @@ def main():
                 st.markdown(message["content"])
 
 
-
 # Calling Main Function
 if __name__ == '__main__':
     main()