Update app.py

app.py

@@ -7,14 +7,23 @@ import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
+from llama_index.core.prompts.prompts import SimpleInputPrompt
+from llama_index.llms.huggingface import HuggingFaceLLM
+from llama_index.legacy.embeddings.langchain import LangchainEmbedding
+from langchain.embeddings.huggingface import HuggingFaceEmbeddings
+from llama_index.core import set_global_service_context, ServiceContext, VectorStoreIndex, Document
+from pathlib import Path
+import fitz  # PyMuPDF
+
 MAX_MAX_NEW_TOKENS = 2048
 DEFAULT_MAX_NEW_TOKENS = 1024
 MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
 
 DESCRIPTION = """\
-# Llama-2 7B Chat
+# Llama-2 7B Chat with Document Context
 
-This Space demonstrates model [Llama-2-7b-chat](https://huggingface.co/meta-llama/Llama-2-7b-chat) by Meta, a Llama 2 model with 7B parameters fine-tuned for chat instructions. Feel free to play with it, or duplicate to run generations without a queue! If you want to run your own service, you can also [deploy the model on Inference Endpoints](https://huggingface.co/inference-endpoints).
+This Space demonstrates model [Llama-2-7b-chat](https://huggingface.co/meta-llama/Llama-2-7b-chat) by Meta, a Llama 2 model with 7B parameters fine-tuned for chat instructions, now enhanced with document-based context. 
+Feel free to play with it, or duplicate to run generations without a queue! If you want to run your own service, you can also [deploy the model on Inference Endpoints](https://huggingface.co/inference-endpoints).
 
 🔎 For more details about the Llama 2 family of models and how to use them with `transformers`, take a look [at our blog post](https://huggingface.co/blog/llama2).
 
@@ -32,13 +41,32 @@ this demo is governed by the original [license](https://huggingface.co/spaces/hu
 if not torch.cuda.is_available():
     DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
 
-
 if torch.cuda.is_available():
-    model_id = "meta-llama/Llama-2-7b-chat-hf"
-    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model_name = "meta-llama/Llama-2-7b-chat-hf"
+    token_file = open("HF_TOKEN.txt")
+    auth_token = token_file.readline().strip()
+
+    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto", token=auth_token)
+    tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir='./model/', token=auth_token)
     tokenizer.use_default_system_prompt = False
 
+    # Load documents and create the index
+    def read_pdf_to_documents(file_path):
+        doc = fitz.open(file_path)
+        documents = []
+        for page_num in range(len(doc)):
+            page = doc.load_page(page_num)
+            text = page.get_text()
+            documents.append(Document(text=text))
+        return documents
+
+    file_path = Path('/content/Full_Pamplet.pdf')  # Update with your document path
+    documents = read_pdf_to_documents(file_path)
+    embeddings = LangchainEmbedding(HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2"))
+    service_context = ServiceContext.from_defaults(chunk_size=1024, embed_model=embeddings)
+    set_global_service_context(service_context)
+    index = VectorStoreIndex.from_documents(documents)
+    query_engine = index.as_query_engine()
 
 @spaces.GPU
 def generate(
@@ -84,63 +112,21 @@ def generate(
         outputs.append(text)
         yield "".join(outputs)
 
+def query_model(question):
+    response = query_engine.query(question)
+    return response.response
 
-chat_interface = gr.ChatInterface(
-    fn=generate,
-    additional_inputs=[
-        gr.Textbox(label="System prompt", lines=6),
-        gr.Slider(
-            label="Max new tokens",
-            minimum=1,
-            maximum=MAX_MAX_NEW_TOKENS,
-            step=1,
-            value=DEFAULT_MAX_NEW_TOKENS,
-        ),
-        gr.Slider(
-            label="Temperature",
-            minimum=0.1,
-            maximum=4.0,
-            step=0.1,
-            value=0.6,
-        ),
-        gr.Slider(
-            label="Top-p (nucleus sampling)",
-            minimum=0.05,
-            maximum=1.0,
-            step=0.05,
-            value=0.9,
-        ),
-        gr.Slider(
-            label="Top-k",
-            minimum=1,
-            maximum=1000,
-            step=1,
-            value=50,
-        ),
-        gr.Slider(
-            label="Repetition penalty",
-            minimum=1.0,
-            maximum=2.0,
-            step=0.05,
-            value=1.2,
-        ),
-    ],
-    stop_btn=None,
-    examples=[
-        ["Hello there! How are you doing?"],
-        ["Can you explain briefly to me what is the Python programming language?"],
-        ["Explain the plot of Cinderella in a sentence."],
-        ["How many hours does it take a man to eat a Helicopter?"],
-        ["Write a 100-word article on 'Benefits of Open-Source in AI research'"],
-    ],
-    cache_examples=False,
+update_prompt_interface = gr.Interface(
+    fn=update_system_prompt,
+    inputs=gr.Textbox(lines=5, placeholder="Enter the system prompt here...", label="System Prompt", value=system_prompt),
+    outputs=gr.Textbox(label="Status"),
+    title="System Prompt Updater",
+    description="Update the system prompt used for context."
 )
 
-with gr.Blocks(css="style.css", fill_height=True) as demo:
-    gr.Markdown(DESCRIPTION)
-    gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button")
-    chat_interface.render()
-    gr.Markdown(LICENSE)
-
-if __name__ == "__main__":
-    demo.queue(max_size=20).launch()
+query_interface = gr.Interface(
+    fn=query_model,
+    inputs=gr.Textbox(lines=2, placeholder="Enter your question here...", label="User Question"),
+    outputs=gr.Textbox(label="Response"),
+    title="Document Query Assistant",
+    description="Ask questions based on the conte