Studymaker2 / app.py
g0th's picture
Update app.py
587fb3d verified
raw
history blame
1.5 kB
import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import PyPDF2
import torch
st.set_page_config(page_title="Perplexity Clone (Gemma)", layout="wide")
st.title("📚 Perplexity-Style AI Study Assistant using Gemma")
# Load Gemma model and tokenizer
@st.cache_resource
def load_model():
tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b-it")
model = AutoModelForCausalLM.from_pretrained(
"google/gemma-7b-it",
torch_dtype=torch.float16,
device_map="auto"
)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=512)
return pipe
textgen = load_model()
# Extract text from uploaded PDF
def extract_text_from_pdf(file):
reader = PyPDF2.PdfReader(file)
text = ""
for page in reader.pages:
text += page.extract_text() + "\n"
return text.strip()
# UI Layout
query = st.text_input("Ask a question or type a query:")
uploaded_file = st.file_uploader("Or upload a PDF to analyze its content:", type=["pdf"])
context = ""
if uploaded_file:
context = extract_text_from_pdf(uploaded_file)
st.text_area("Extracted Content", context, height=200)
if st.button("Generate Answer"):
with st.spinner("Generating with Gemma..."):
prompt = query
if context:
prompt = f"Context:\n{context}\n\nQuestion: {query}"
output = textgen(prompt)[0]["generated_text"]
st.success("Answer:")
st.write(output)