Update app.py

app.py

@@ -1,15 +1,157 @@
-import streamlit as st
+# import streamlit as st
+# from ctransformers import AutoModelForCausalLM
+# from fastapi import FastAPI
+# from pydantic import BaseModel
+# import time
+# from sklearn.feature_extraction.text import TfidfVectorizer
+# from sklearn.metrics.pairwise import cosine_similarity
+# import numpy as np
+# import PyPDF2
+
+# # Load the model
+# @st.cache_resource
+# def load_model():
+#     return AutoModelForCausalLM.from_pretrained(
+#         "llama-2-7b-chat.Q5_K_M.gguf",
+#         model_type='llama',
+#         max_new_tokens=1096,
+#         threads=3,
+#     )
+
+# llm = load_model()
+
+# # Pydantic object
+# class Validation(BaseModel):
+#     prompt: str
+#     context: str
+
+# # FastAPI
+# app = FastAPI()
+
+# @app.post("/llm_on_cpu")
+# async def stream(item: Validation):
+#     system_prompt = "You are an intelligent assistant. Answer the question based solely on the provided context. If the answer cannot be found in the context, say 'I don't have enough information to answer this question based on the provided context.'"
+#     B_INST, E_INST = "[INST]", "[/INST]"
+#     B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
+    
+#     full_prompt = f"{B_INST}{B_SYS}{system_prompt}{E_SYS}Context: {item.context}\n\nQuestion: {item.prompt}\n\nAnswer:{E_INST}"
+    
+#     return llm(full_prompt, stream=True)
+
+# # Helper functions
+# @st.cache_resource
+# def get_vectorizer():
+#     return TfidfVectorizer()
+
+# def encode_texts(texts, vectorizer):
+#     return vectorizer.fit_transform(texts)
+
+# def retrieve_documents(query, documents, top_k=3):
+#     vectorizer = get_vectorizer()
+#     all_texts = documents + [query]
+#     embeddings = encode_texts(all_texts, vectorizer)
+#     doc_embeddings = embeddings[:-1]
+#     query_embedding = embeddings[-1]
+#     similarities = cosine_similarity(query_embedding, doc_embeddings).flatten()
+#     top_indices = np.argsort(similarities)[-top_k:][::-1]
+#     return [documents[i] for i in top_indices]
+
+# def generate_llama_response(prompt, context):
+#     system_prompt = "You are an intelligent assistant. Answer the question based solely on the provided context. If the answer cannot be found in the context, say 'I don't have enough information to answer this question based on the provided context.'"
+#     B_INST, E_INST = "[INST]", "[/INST]"
+#     B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
+    
+#     full_prompt = f"{B_INST}{B_SYS}{system_prompt}{E_SYS}Context: {context}\n\nQuestion: {prompt}\n\nAnswer:{E_INST}"
+    
+#     return llm(full_prompt, stream=True)
+
+# def stream_response(prompt, context, placeholder):
+#     response = ""
+#     for token in generate_llama_response(prompt, context):
+#         response += token
+#         placeholder.markdown(response + "▌")
+#     placeholder.markdown(response)
+
+# def summarize_pdf(documents, placeholder):
+#     context = " ".join(documents[:3])
+#     prompt = "Provide a concise summary of the main points in the following text:"
+#     stream_response(prompt, context, placeholder)
+
+# def create_questions(documents, placeholder):
+#     context = " ".join(documents[:3])
+#     prompt = "Generate 5 important questions based on the following text. These questions should cover key concepts and encourage critical thinking:"
+#     stream_response(prompt, context, placeholder)
+
+# def create_notes(documents, placeholder):
+#     context = " ".join(documents[:3])
+#     prompt = "Create concise and informative notes based on the following text. Focus on key concepts, definitions, and important points:"
+#     stream_response(prompt, context, placeholder)
+
+# def extract_text_from_pdf(file):
+#     pdf_reader = PyPDF2.PdfReader(file)
+#     text = ""
+#     for page in pdf_reader.pages:
+#         text += page.extract_text()
+#     return text
+
+# # Streamlit UI
+# def main():
+#     st.title("LLaMA-2 PDF Analysis and Chat Interface")
+
+#     uploaded_file = st.file_uploader("Choose a PDF file", type="pdf")
+
+#     if uploaded_file is not None:
+#         text = extract_text_from_pdf(uploaded_file)
+#         chunk_size = 1000
+#         documents = [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]
+
+#         st.success("PDF uploaded and processed successfully!")
+
+#         # Create columns for buttons
+#         col1, col2, col3 = st.columns(3)
+
+#         with col1:
+#             if st.button("Summarize PDF"):
+#                 with st.spinner("Generating summary..."):
+#                     placeholder = st.empty()
+#                     summarize_pdf(documents, placeholder)
+
+#         with col2:
+#             if st.button("Generate Questions"):
+#                 with st.spinner("Generating questions..."):
+#                     placeholder = st.empty()
+#                     create_questions(documents, placeholder)
+
+#         with col3:
+#             if st.button("Create Notes"):
+#                 with st.spinner("Creating notes..."):
+#                     placeholder = st.empty()
+#                     create_notes(documents, placeholder)
+
+#         # Chat interface
+#         st.subheader("Chat with the PDF")
+#         user_question = st.text_input("Ask a question about the PDF:")
+#         if user_question:
+#             relevant_docs = retrieve_documents(user_question, documents)
+#             context = " ".join(relevant_docs)
+#             with st.spinner("Generating answer..."):
+#                 placeholder = st.empty()
+#                 stream_response(user_question, context, placeholder)
+
+#     else:
+#         st.info("Please upload a PDF file to begin.")
+
+# if __name__ == "__main__":
+#     main()
+
+import gradio as gr
 from ctransformers import AutoModelForCausalLM
-from fastapi import FastAPI
-from pydantic import BaseModel
-import time
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.metrics.pairwise import cosine_similarity
 import numpy as np
 import PyPDF2
 
 # Load the model
-@st.cache_resource
 def load_model():
     return AutoModelForCausalLM.from_pretrained(
         "llama-2-7b-chat.Q5_K_M.gguf",
@@ -20,26 +162,7 @@ def load_model():
 
 llm = load_model()
 
-# Pydantic object
-class Validation(BaseModel):
-    prompt: str
-    context: str
-
-# FastAPI
-app = FastAPI()
-
-@app.post("/llm_on_cpu")
-async def stream(item: Validation):
-    system_prompt = "You are an intelligent assistant. Answer the question based solely on the provided context. If the answer cannot be found in the context, say 'I don't have enough information to answer this question based on the provided context.'"
-    B_INST, E_INST = "[INST]", "[/INST]"
-    B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
-    
-    full_prompt = f"{B_INST}{B_SYS}{system_prompt}{E_SYS}Context: {item.context}\n\nQuestion: {item.prompt}\n\nAnswer:{E_INST}"
-    
-    return llm(full_prompt, stream=True)
-
 # Helper functions
-@st.cache_resource
 def get_vectorizer():
     return TfidfVectorizer()
 
@@ -63,83 +186,76 @@ def generate_llama_response(prompt, context):
     
     full_prompt = f"{B_INST}{B_SYS}{system_prompt}{E_SYS}Context: {context}\n\nQuestion: {prompt}\n\nAnswer:{E_INST}"
     
-    return llm(full_prompt, stream=True)
+    return llm(full_prompt)
+
+def extract_text_from_pdf(file):
+    pdf_reader = PyPDF2.PdfReader(file)
+    text = ""
+    for page in pdf_reader.pages:
+        text += page.extract_text()
+    return text
 
-def stream_response(prompt, context, placeholder):
-    response = ""
-    for token in generate_llama_response(prompt, context):
-        response += token
-        placeholder.markdown(response + "▌")
-    placeholder.markdown(response)
+# Gradio functions
+def process_pdf(file):
+    text = extract_text_from_pdf(file)
+    chunk_size = 1000
+    documents = [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]
+    return documents
 
-def summarize_pdf(documents, placeholder):
+def summarize_pdf(documents):
     context = " ".join(documents[:3])
     prompt = "Provide a concise summary of the main points in the following text:"
-    stream_response(prompt, context, placeholder)
+    return generate_llama_response(prompt, context)
 
-def create_questions(documents, placeholder):
+def create_questions(documents):
     context = " ".join(documents[:3])
     prompt = "Generate 5 important questions based on the following text. These questions should cover key concepts and encourage critical thinking:"
-    stream_response(prompt, context, placeholder)
+    return generate_llama_response(prompt, context)
 
-def create_notes(documents, placeholder):
+def create_notes(documents):
     context = " ".join(documents[:3])
     prompt = "Create concise and informative notes based on the following text. Focus on key concepts, definitions, and important points:"
-    stream_response(prompt, context, placeholder)
+    return generate_llama_response(prompt, context)
 
-def extract_text_from_pdf(file):
-    pdf_reader = PyPDF2.PdfReader(file)
-    text = ""
-    for page in pdf_reader.pages:
-        text += page.extract_text()
-    return text
+def chat_with_pdf(documents, user_question):
+    relevant_docs = retrieve_documents(user_question, documents)
+    context = " ".join(relevant_docs)
+    return generate_llama_response(user_question, context)
+
+# Gradio interface
+def gradio_interface(pdf_file):
+    documents = process_pdf(pdf_file)
+    
+    summary = summarize_pdf(documents)
+    questions = create_questions(documents)
+    notes = create_notes(documents)
+    
+    def chat(user_question):
+        return chat_with_pdf(documents, user_question)
+    
+    return (
+        gr.Textbox(value=summary, label="PDF Summary"),
+        gr.Textbox(value=questions, label="Generated Questions"),
+        gr.Textbox(value=notes, label="Created Notes"),
+        gr.Textbox(placeholder="Ask a question about the PDF", label="Chat with PDF"),
+        gr.Textbox(label="Answer"),
+        gr.Button("Ask"),
+    )
+
+iface = gr.Interface(
+    fn=gradio_interface,
+    inputs=gr.File(label="Upload PDF"),
+    outputs=[
+        gr.Textbox(label="PDF Summary"),
+        gr.Textbox(label="Generated Questions"),
+        gr.Textbox(label="Created Notes"),
+        gr.Textbox(label="Chat with PDF"),
+        gr.Textbox(label="Answer"),
+        gr.Button("Ask"),
+    ],
+    title="LLaMA-2 PDF Analysis and Chat Interface",
+    description="Upload a PDF file to analyze and chat with its content using LLaMA-2.",
+    allow_flagging="never",
+)
 
-# Streamlit UI
-def main():
-    st.title("LLaMA-2 PDF Analysis and Chat Interface")
-
-    uploaded_file = st.file_uploader("Choose a PDF file", type="pdf")
-
-    if uploaded_file is not None:
-        text = extract_text_from_pdf(uploaded_file)
-        chunk_size = 1000
-        documents = [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]
-
-        st.success("PDF uploaded and processed successfully!")
-
-        # Create columns for buttons
-        col1, col2, col3 = st.columns(3)
-
-        with col1:
-            if st.button("Summarize PDF"):
-                with st.spinner("Generating summary..."):
-                    placeholder = st.empty()
-                    summarize_pdf(documents, placeholder)
-
-        with col2:
-            if st.button("Generate Questions"):
-                with st.spinner("Generating questions..."):
-                    placeholder = st.empty()
-                    create_questions(documents, placeholder)
-
-        with col3:
-            if st.button("Create Notes"):
-                with st.spinner("Creating notes..."):
-                    placeholder = st.empty()
-                    create_notes(documents, placeholder)
-
-        # Chat interface
-        st.subheader("Chat with the PDF")
-        user_question = st.text_input("Ask a question about the PDF:")
-        if user_question:
-            relevant_docs = retrieve_documents(user_question, documents)
-            context = " ".join(relevant_docs)
-            with st.spinner("Generating answer..."):
-                placeholder = st.empty()
-                stream_response(user_question, context, placeholder)
-
-    else:
-        st.info("Please upload a PDF file to begin.")
-
-if __name__ == "__main__":
-    main()
+iface.launch()