inam09 commited on
Commit
3c31d7d
·
verified ·
1 Parent(s): f03d0ce

Update src/streamlit_app.py

Browse files
Files changed (1) hide show
  1. src/streamlit_app.py +7 -2
src/streamlit_app.py CHANGED
@@ -33,10 +33,15 @@ nltk.download('punkt_tab', download_dir=nltk_data_dir, quiet=True)
33
 
34
 
35
  # Load models
36
- summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", device=torch.device("cpu"))
 
 
 
 
 
 
37
  qa_pipeline = pipeline("question-answering", model="distilbert-base-cased-distilled-squad", device=torch.device("cpu"))
38
  embedding_model = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
39
-
40
  vector_dbs = {} # Dictionary to store multiple vector databases, keyed by document title
41
  extracted_texts = {} # Dictionary to store extracted text, keyed by document title
42
  current_doc_title = None
 
33
 
34
 
35
  # Load models
36
+ # summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", device=torch.device("cpu"))
37
+ summarizer = pipeline(
38
+ "summarization",
39
+ model="sshleifer/distilbart-cnn-12-6",
40
+ from_flax=True,
41
+ device=-1 # CPU mode, use device=0 for GPU
42
+ )
43
  qa_pipeline = pipeline("question-answering", model="distilbert-base-cased-distilled-squad", device=torch.device("cpu"))
44
  embedding_model = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
 
45
  vector_dbs = {} # Dictionary to store multiple vector databases, keyed by document title
46
  extracted_texts = {} # Dictionary to store extracted text, keyed by document title
47
  current_doc_title = None