Update app.py

app.py

@@ -1,19 +1,18 @@
 from langchain_core.prompts import PromptTemplate
 from langchain.chains import create_retrieval_chain
 from langchain.chains.combine_documents import create_stuff_documents_chain
-import gradio as gr
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+# import gradio as gr
 import numpy as np
 from langchain_ollama import OllamaLLM
 from langchain_huggingface import HuggingFaceEmbeddings
-from langchain_community.llms import HuggingFacePipeline
+# from langchain_community.llms import HuggingFacePipeline
 from load_document import load_data
 from split_document import split_docs
 from embed_docs import embed_docs
 from retrieve import retrieve
 from datetime import datetime
-# from js import js
-# from theme import theme
+from js import js
+from theme import theme
 import os
 import glob
 from fastapi import FastAPI, Query, Request
@@ -40,34 +39,13 @@ embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6
 
 def fetch_doc():
     # Adjust the path as needed, e.g., './' for current directory
-    pdf_files = glob.glob("Document/*.pdf")
-
-    # If you want to include subdirectories:
-    # pdf_files = glob.glob("**/*.pdf", recursive=True)
+    pdf_files = glob.glob("*.pdf")
 
     return pdf_files
 
 # # Define llm
 hf_token = os.environ.get("HF_TOKEN").strip()  # Ensure to set your Hugging Face token in the environment variable HF_TOKEN
-# #llm = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.3", device="cpu", use_auth_token=hf_token, token=hf_token)
-# #llm = OllamaLLM(model="mistral:7b-instruct", base_url="http://host.docker.internal:11434")
-model_id = "google/gemma-2b-it"
-
-# # Load tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
-model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cpu", torch_dtype="auto", token=hf_token)
-
-# # Create text generation pipeline
-hf_pipe = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_new_tokens=512,
-    temperature=0.7,
-    top_p=0.9,
-    do_sample=True
-)
-llm = HuggingFacePipeline(pipeline=hf_pipe)
+llm = OllamaLLM(model="mistral:7b-instruct")
 
 pdf_files = fetch_doc() #Fetch Dataset
 chunks = None