File size: 4,002 Bytes
4b09806 | 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 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | import streamlit as st
from PyPDF2 import PdfReader
from transformers import AutoTokenizer, AutoModel
import torch
import faiss
import numpy as np
import os
import requests
# -------------------- Config --------------------
GROQ_API_KEY = os.getenv("GROQ_API_KEY") # Add this as a secret in Hugging Face
GROQ_API_URL = "https://api.groq.com/openai/v1/chat/completions"
EMBED_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
# -------------------- PDF Processing --------------------
def read_pdf(file):
pdf = PdfReader(file)
text = ""
for page in pdf.pages:
text += page.extract_text() + "\n"
return text
# -------------------- Chunking --------------------
def chunk_text(text, chunk_size=500, overlap=50):
words = text.split()
chunks = []
start = 0
while start < len(words):
end = start + chunk_size
chunk = " ".join(words[start:end])
chunks.append(chunk)
start += chunk_size - overlap
return chunks
# -------------------- Embedding --------------------
@st.cache_resource
def load_embedding_model():
tokenizer = AutoTokenizer.from_pretrained(EMBED_MODEL_NAME)
model = AutoModel.from_pretrained(EMBED_MODEL_NAME)
return tokenizer, model
def get_embeddings(text_chunks, tokenizer, model):
embeddings = []
for chunk in text_chunks:
inputs = tokenizer(chunk, return_tensors="pt", truncation=True, padding=True)
with torch.no_grad():
outputs = model(**inputs)
emb = outputs.last_hidden_state[:, 0, :].numpy()[0]
embeddings.append(emb)
return np.array(embeddings)
# -------------------- FAISS --------------------
def build_faiss_index(embeddings):
dimension = embeddings.shape[1]
index = faiss.IndexFlatL2(dimension)
index.add(embeddings)
return index
def search_index(index, query, tokenizer, model, chunks, top_k=3):
inputs = tokenizer(query, return_tensors="pt", truncation=True, padding=True)
with torch.no_grad():
outputs = model(**inputs)
query_emb = outputs.last_hidden_state[:, 0, :].numpy()
distances, indices = index.search(query_emb, top_k)
return [chunks[i] for i in indices[0]]
# -------------------- GROQ Query --------------------
def query_groq(context, question):
prompt = f"""You are a helpful engineering assistant. Use the following context to answer the question.
Context:
{context}
Question:
{question}
"""
headers = {
"Authorization": f"Bearer {GROQ_API_KEY}",
"Content-Type": "application/json"
}
payload = {
"model": "llama3-8b-8192",
"messages": [
{"role": "system", "content": "You are a helpful engineering tutor."},
{"role": "user", "content": prompt}
],
"temperature": 0.3,
"max_tokens": 512
}
response = requests.post(GROQ_API_URL, headers=headers, json=payload)
response.raise_for_status()
return response.json()["choices"][0]["message"]["content"]
# -------------------- Streamlit UI --------------------
st.title("📚 engGlass RAG Assistant")
st.write("Upload a PDF, ask engineering questions, and get smart answers!")
uploaded_file = st.file_uploader("Upload PDF", type="pdf")
question = st.text_input("Ask a question based on the uploaded document:")
if uploaded_file and question:
with st.spinner("Reading and processing PDF..."):
text = read_pdf(uploaded_file)
chunks = chunk_text(text)
tokenizer, model = load_embedding_model()
embeddings = get_embeddings(chunks, tokenizer, model)
index = build_faiss_index(embeddings)
top_chunks = search_index(index, question, tokenizer, model, chunks)
context = "\n".join(top_chunks)
with st.spinner("Generating answer from Groq..."):
try:
answer = query_groq(context, question)
st.markdown("### 💡 Answer")
st.write(answer)
except Exception as e:
st.error(f"Error: {str(e)}")
|