DVC-AI-MODEL / app.py
Emalawi19's picture
Update app.py
81f4994 verified
Raw
History Blame Contribute Delete
1.57 kB
import gradio as gr
import PyPDF2
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
# Load models
embed_model = SentenceTransformer('all-MiniLM-L6-v2')
model_name = "google/flan-t5-base"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
# Load PDF
pdf_file = "FA1.pdf"
reader = PyPDF2.PdfReader(pdf_file)
text = ""
for page in reader.pages:
text += page.extract_text() + "\n"
# Chunk text
chunk_size = 500
chunks = [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]
# Create embeddings
embeddings = embed_model.encode(chunks)
index = faiss.IndexFlatL2(embeddings.shape[1])
index.add(np.array(embeddings).astype('float32'))
# Search function
def search(query, k=3):
q_emb = embed_model.encode([query]).astype('float32')
_, indices = index.search(q_emb, k)
return [chunks[i] for i in indices[0]]
# Chat function
def chat(question):
context = "\n".join(search(question))
prompt = f"""
Answer ONLY from context.
If not found say: Not found in document.
Context:
{context}
Question:
{question}
Answer:
"""
inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
outputs = model.generate(**inputs, max_new_tokens=120)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
# Gradio UI
interface = gr.Interface(
fn=chat,
inputs="text",
outputs="text",
title="Emalawi19 AI Assistant"
)
interface.launch()