Naznazrin/titleabstractrag

1import pandas as pd2import numpy as np3import torch4from sentence_transformers import SentenceTransformer5from transformers import AutoTokenizer, AutoModelForSeq2SeqLM6 7# Load models (cached)8@st.cache_resource9def load_models():10    embedding_model = SentenceTransformer('all-MiniLM-L6-v2')11    tokenizer = AutoTokenizer.from_pretrained('google/flan-t5-small')12    model = AutoModelForSeq2SeqLM.from_pretrained('google/flan-t5-small')13    device = 'cuda' if torch.cuda.is_available() else 'cpu'14    model.to(device)15    return embedding_model, tokenizer, model, device16 17embedding_model, tokenizer_rag, model_rag, device = load_models()18 19# App title20st.set_page_config(page_title="RAG-Based Title & Abstract Screening", layout="wide")21st.title("📰 RAG-Based Title & Abstract Screening")22 23# Sidebar: file upload and query24st.sidebar.header("1. Upload Dataset")25uploaded_file = st.sidebar.file_uploader("Upload an Excel file", type=["xlsx"])26 27st.sidebar.header("2. Define Screening Query & Settings")28query = st.sidebar.text_input("Enter your screening query:")29k = st.sidebar.slider("Number of papers to retrieve", 1, 10, 5)30 31if uploaded_file and query:32    # Load data33    df = pd.read_excel(uploaded_file)34    if 'Title' not in df.columns or 'Abstract' not in df.columns:35        st.error("Excel must contain 'Title' and 'Abstract' columns.")36    else:37        # Prepare corpus & FAISS38        df = df.dropna(subset=['Title', 'Abstract']).reset_index(drop=True)39        corpus = (df['Title'] + " " + df['Abstract']).tolist()40        embeddings = embedding_model.encode(corpus, convert_to_numpy=True)41        dim = embeddings.shape[1]42        index = faiss.IndexFlatL2(dim)43        index.add(embeddings)44 45        # Retrieval46        query_emb = embedding_model.encode(query, convert_to_numpy=True).reshape(1, -1)47        _, indices = index.search(np.array([query_emb]), k)48        retrieved_idxs = indices[0]49 50        # Generation51        context = " ".join(corpus[i] for i in retrieved_idxs)52        prompt = f"Based on the following papers, decide whether to include or exclude them for the query: '{query}'. Just return 'Include' or 'Exclude'. Context: {context}"53        inputs = tokenizer_rag(prompt, return_tensors='pt', truncation=True, max_length=1024).to(device)54        outputs = model_rag.generate(**inputs, max_new_tokens=10)55        decision = tokenizer_rag.decode(outputs[0], skip_special_tokens=True).strip()56 57        # Display results58        st.header("🔍 Screening Decision")59        st.subheader(decision)60 61        st.header(f"📄 Top {k} Retrieved Papers")62        results = []63        for idx in retrieved_idxs:64            title = df.loc[idx, 'Title']65            abstract = df.loc[idx, 'Abstract']66            st.markdown(f"**Title:** {title}")67            st.markdown(f"**Abstract:** {abstract}")68            st.markdown(f"**Prediction:** {decision}")69            st.markdown("---")70            results.append({71                'Title': title,72                'Abstract': abstract,73                'Prediction': decision74            })75 76        # Download results77        result_df = pd.DataFrame(results)78        st.download_button(79            label="Download results as Excel",80            data=result_df.to_excel(index=False),81            file_name="rag_screening_results.xlsx",82            mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"83        )84else:85    if not uploaded_file:86        st.info("Please upload an Excel file to begin.")87    elif not query:88        st.info("Please enter a screening query.")89