Update src/streamlit_app.py

src/streamlit_app.py

@@ -5,7 +5,7 @@ from langchain.text_splitter import CharacterTextSplitter
 from langchain.schema import Document
 from langchain.embeddings import SentenceTransformerEmbeddings
 from langchain.vectorstores import FAISS
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForQuestionAnswering, pipeline
 
 # Cache the QA initialization so ingestion runs once per session
 @st.cache_resource
@@ -37,28 +37,27 @@ def init_qa(zip_bytes):
     embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
     vector_store = FAISS.from_documents(split_docs, embeddings)
 
-    # Load the RAG model
-    generator = pipeline(
-        "text2text-generation",
-        model="PleIAs/Pleias-RAG-350M",
-        tokenizer="PleIAs/Pleias-RAG-350M"
-    )
+    # Load the QA model and tokenizer
+    tokenizer = AutoTokenizer.from_pretrained("deepset/roberta-base-squad2")
+    model = AutoModelForQuestionAnswering.from_pretrained("deepset/roberta-base-squad2")
+    qa_pipeline = pipeline("question-answering", model=model, tokenizer=tokenizer)
 
-    return vector_store, generator
+    return vector_store, qa_pipeline
 
 # Streamlit UI
-st.title("Pleias-RAG 350M Streamlit App")
-st.write("Upload a ZIP of PDFs to initialize the RAG engine.")
+st.title("RoBERTa QA Streamlit App")
+st.write("Upload a ZIP of PDFs to initialize the QA engine.")
 zip_file = st.file_uploader("ZIP file", type=["zip"])
 
 if zip_file:
-    vector_store, generator = init_qa(zip_file.read())
+    vector_store, qa = init_qa(zip_file.read())
     query = st.text_input("Ask a question:")
     if query:
         docs = vector_store.similarity_search(query, k=4)
         context = "\n\n".join([doc.page_content for doc in docs])
-        prompt = f"question: {query}\ncontext: {context}"
-        answer = generator(prompt, max_length=512, do_sample=False)[0]["generated_text"]
+        # Run QA
+        result = qa(question=query, context=context)
+        answer = result.get("answer", "No answer found.")
         st.write(answer)
 else:
     st.info("Awaiting ZIP upload.")