Update app.py

app.py

@@ -8,7 +8,7 @@ from bs4 import BeautifulSoup
 from sentence_transformers import SentenceTransformer
 from langchain_community.vectorstores import FAISS
 from langchain.embeddings.base import Embeddings
-import google.generativeai as genai
+from transformers import pipeline
 
 # === Embeddings Wrapper ===
 class SentenceTransformerEmbeddings(Embeddings):
@@ -35,27 +35,28 @@ def split_text(text, chunk_size=500, overlap=50):
         start += chunk_size - overlap
     return chunks
 
-def ask_gemini(question, context, api_key):
-    genai.configure(api_key=api_key)
-    model = genai.GenerativeModel("gemini-pro")
-    prompt = f"""You are a helpful assistant. Use the context below to answer the question.
-
-Context:
-{context}
-
-Question: {question}
-Answer:"""
-    response = model.generate_content(prompt)
-    return response.text
+def ask_mistral(question, context, hf_api_key):
+    # Load the Hugging Face Mistral model pipeline
+    nlp = pipeline("question-answering", model="mistral-7b", tokenizer="mistral-7b", use_auth_token=hf_api_key)
+    
+    # Format the input
+    inputs = {
+        'context': context,
+        'question': question
+    }
+    
+    # Generate the answer using Mistral
+    answer = nlp(inputs)
+    return answer['answer']
 
 def create_vectorstore(chunks):
     embeddings = SentenceTransformerEmbeddings()
     return FAISS.from_texts(chunks, embedding=embeddings)
 
-def generate_answer(vectorstore, question, api_key):
+def generate_answer(vectorstore, question, hf_api_key):
     docs = vectorstore.similarity_search(question, k=3)
     context = "\n".join([doc.page_content for doc in docs])
-    return ask_gemini(question, context, api_key), docs
+    return ask_mistral(question, context, hf_api_key), docs
 
 def extract_website_text(url):
     try:
@@ -75,7 +76,7 @@ st.title("🔍 RAG Assistant: Chat with PDF, CSV, or Website")
 # Sidebar
 with st.sidebar:
     data_source = st.selectbox("📂 Select Input Type", ["PDF", "CSV", "Website URL"])
-    gemini_api_key = st.text_input("🔑 Enter Gemini API Key", type="password")
+    hf_api_key = st.text_input("🔑 Enter Hugging Face API Key", type="password")
 
 # === Logic by Data Source ===
 vectorstore = None
@@ -120,12 +121,12 @@ elif data_source == "Website URL":
             st.success("✅ Website text extracted and indexed!")
 
 # === QA Section ===
-if vectorstore and gemini_api_key:
+if vectorstore and hf_api_key:
     st.subheader("❓ Ask a Question")
     question = st.text_input("💬 Your question")
     if question:
         with st.spinner("🔍 Thinking..."):
-            answer, top_docs = generate_answer(vectorstore, question, gemini_api_key)
+            answer, top_docs = generate_answer(vectorstore, question, hf_api_key)
             st.success("🧠 Answer")
             st.write(answer)
 
@@ -135,6 +136,5 @@ if vectorstore and gemini_api_key:
 
             st.download_button("📤 Download Answer", answer, file_name="rag_answer.txt")
 
-elif not gemini_api_key:
-    st.info("🔐 Please enter your Gemini API key in the sidebar.")
-
+elif not hf_api_key:
+    st.info("🔐 Please enter your Hugging Face API key in the sidebar.")