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()