File size: 2,405 Bytes
6459544 02cc2b1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
import streamlit as st
import os
import openai
from sentence_transformers import SentenceTransformer, util
import PyPDF2
openai.api_key = os.getenv("openapikey")
model = SentenceTransformer('all-MiniLM-L6-v2')
def load_pdf(uploaded_file):
with open("temp.pdf", "wb") as f:
f.write(uploaded_file.getvalue())
with open("temp.pdf", 'rb') as file:
pdf_reader = PyPDF2.PdfReader(file)
text = ""
for page_num in range(len(pdf_reader.pages)):
page = pdf_reader.pages[page_num]
text += page.extract_text()
return text
def chunk_text(text, chunk_size=500, overlap=100):
chunks = []
for i in range(0, len(text), chunk_size - overlap):
chunks.append(text[i:i + chunk_size])
return chunks
def create_embeddings(chunks):
embeddings = model.encode(chunks, convert_to_tensor=True)
return embeddings
def find_relevant_chunks(query_embedding, chunk_embeddings, chunks, top_k=3):
cosine_scores = util.pytorch_cos_sim(query_embedding, chunk_embeddings)[0]
top_results = sorted(range(len(cosine_scores)), key=lambda i: cosine_scores[i], reverse=True)[:top_k]
relevant_chunks = [chunks[i] for i in top_results]
return relevant_chunks
def generate_response(query, context):
messages = [
{"role": "system", "content": "You are a helpful assistant that answers questions based on provided context."},
{"role": "user", "content": f"Context: {context}\nQuestion: {query}"}
]
response = openai.chat.completions.create(
model="gpt-3.5-turbo", # Or another suitable chat model
messages=messages,
max_tokens=200
)
return response.choices[0].message.content.strip()
st.title("Simple RAG Application (No LangChain)")
uploaded_file = st.file_uploader("Upload PDF", type="pdf")
if uploaded_file:
with st.spinner("Processing PDF..."):
pdf_text = load_pdf(uploaded_file)
chunks = chunk_text(pdf_text)
chunk_embeddings = create_embeddings(chunks)
query = st.text_input("Ask a question:")
if query:
query_embedding = model.encode([query], convert_to_tensor=True)
relevant_chunks = find_relevant_chunks(query_embedding, chunk_embeddings, chunks)
context = "\n".join(relevant_chunks)
answer = generate_response(query, context)
st.write("Answer:", answer) |