VirtualTrainR / app.py
aimanathar's picture
Create app.py
e6365ab verified
import torch, textwrap, numpy as np, faiss
import warnings
warnings.filterwarnings("ignore")
from pdfminer.high_level import extract_text
from pdf2image import convert_from_path
import pytesseract
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import gradio as gr
# ================== PDF Handling Functions ==================
def pdf_to_text(path):
try:
txt = extract_text(path) or ""
except Exception:
txt = ""
if len(txt.strip()) < 200:
try:
pages = convert_from_path(path, dpi=200)
ocr_all = [pytesseract.image_to_string(img) for img in pages]
txt = "\n".join(ocr_all)
except Exception:
txt = ""
return txt
def chunk_text(text, max_chars=800):
paras = [p.strip() for p in text.split("\n") if p.strip()]
chunks, buf = [], ""
for p in paras:
if len(p) > max_chars:
for piece in textwrap.wrap(p, width=max_chars, break_long_words=False):
chunks.append(piece.strip())
else:
if len(buf) + len(p) + 1 <= max_chars:
buf = (buf + "\n" + p).strip()
else:
if buf: chunks.append(buf)
buf = p
if buf: chunks.append(buf)
return [c for c in chunks if len(c) > 80]
# ================== Load Embeddings + Model ==================
embed_model = SentenceTransformer("all-MiniLM-L6-v2")
model_id = "google/flan-t5-base"
tok = AutoTokenizer.from_pretrained(model_id)
gen_model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
device = "cuda" if torch.cuda.is_available() else "cpu"
gen_model.to(device)
# ================== Chat Function ==================
def chat_fn(message, history=None):
prompt = f"Answer clearly and exam-ready:\n\nQuestion:\n{message}"
inputs = tok(prompt, return_tensors="pt", truncation=True, padding=True, max_length=1024).to(device)
out = gen_model.generate(**inputs, max_new_tokens=120, num_beams=4, do_sample=False)
return tok.decode(out[0], skip_special_tokens=True).strip()
# ================== Gradio Interface ==================
iface = gr.ChatInterface(
fn=chat_fn,
title="💬 Practical Chatbot",
description="Ask about Physics & Chemistry Practicals (Class 9–10)."
)
iface.launch()