app.py

import os
import streamlit as st
from get_pat_data import Patent_DataCreator
from datasets import load_dataset
import re
import boto3
import time
import requests
from bs4 import BeautifulSoup
import pandas as pd
import pinecone
import torch.nn.functional as F
from torch import Tensor
from transformers import AutoTokenizer, AutoModel
from keybert import KeyBERT
from keyphrase_vectorizers import KeyphraseCountVectorizer
kw_model=KeyBERT(model='AI-Growth-Lab/PatentSBERTa')

s3 = boto3.resource('s3',
     region_name='us-east-1',
     aws_access_key_id='AKIA3VGKPNV5NSVBJWEE',
     aws_secret_access_key='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')

INDEX_API_KEY='b33ddf5d-5b1a-4d0e-9a3f-572008563791'
INDEX_DIMENSION=768
INDEX_ENV='gcp-starter'
INDEX_NAME='wiki-index'

# getting Pinecone credntials
# INDEX_DIMENSION=768
# logging.info(f"Index dimensions are:{INDEX_DIMENSION}")
pinecone.init(api_key=INDEX_API_KEY, environment=INDEX_ENV)
index = pinecone.Index(index_name=INDEX_NAME )

tokenizer = AutoTokenizer.from_pretrained('intfloat/e5-base')
model = AutoModel.from_pretrained('intfloat/e5-base')

# data=pd.read_csv("wikicat_all.csv")
         
def get_pat_text(pnkc_no):
    pat_data=Patent_DataCreator(pnkc_no)
    bib_key,pnkc_without_kindcode,pnkc_suffix=pat_data.get_bib_key()
    bib_bucket=pat_data.get_bib_bucket()
    bib_data=pat_data.get_bib_data(s3)
    claims_data=pat_data.get_claims_data(s3)
    desc_data=pat_data.get_desc_data(s3)
    df1,df2,df3=pat_data.get_patent_dfs()
    dataset=pat_data.get_patent_dataset()

    Title=dataset[1]['Title'][0]
    Abstract=dataset[1]['Abstract'][0]
    Claims=dataset[1]['Claims'][0]
    Description=dataset[1]['Description'][0]
    # SOI=dataset[1]['SOI'][0]

    pat_text= Title+Abstract
    return pat_text


# Function to fetch categories, title, and related text from a Wikipedia page
def fetch_wikipedia_data(article_title):
    url = f"https://en.wikipedia.org/wiki/{article_title.replace(' ', '_')}"

    response = requests.get(url)

    if response.status_code == 200:
        soup = BeautifulSoup(response.text, 'html.parser')

        # Find the categories section at the bottom of the page
        categories_section = soup.find("div", {"class": "mw-normal-catlinks"})

        if categories_section:
            # Extract individual categories
            categories = [cat.text for cat in categories_section.find("ul").find_all("li")]

            # Extract the title
            title = article_title

            return {"title": title, "categories": categories}

    return None

def get_wiki_category_aprch_1(pat_text):
    print(pat_text)
    keywords=kw_model.extract_keywords(pat_text,keyphrase_ngram_range=(1, 3),top_n=15,vectorizer=KeyphraseCountVectorizer())
    titles=[]
    for i in range(len(keywords)):
        title=keywords[i][0]
        titles.append(title)
    data = []
    for i in titles:
        results = fetch_wikipedia_data(i)
        data.append(results)
    cats=[]
    for i in range(len(data)):
        if data[i] is not None:
            cat=data[i]['categories']
            cats.append(cat)
    result=[j for i in cats for j in i]
    res = [i for n, i in enumerate(result) if i not in result[:n]]
    return titles,res

# def get_wiki_category_aprch_2(pat_text):
#     print(pat_text)
#     keywords=kw_model.extract_keywords(pat_text,keyphrase_ngram_range=(1, 3),top_n=10,vectorizer=KeyphraseCountVectorizer())
#     titles=[]
#     for i in range(len(keywords)):
#         title=keywords[i][0]
#         titles.append(title)
#     data = []
#     for i in titles:
#         results = fetch_wikipedia_data(i)
#         data.append(results)
#     cats=[]
#     for i in range(len(data)):
#         if data[i] is not None:
#             cat=data[i]['categories']
#             cats.append(cat)
#     result=[j for i in cats for j in i]
#     res = [i for n, i in enumerate(result) if i not in result[:n]]
#     return res

def average_pool(last_hidden_states: Tensor,
                attention_mask: Tensor) -> Tensor:
    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
            
def get_wiki_category(pat_text):
    # print(pat_text)
    keywords=kw_model.extract_keywords(pat_text,keyphrase_ngram_range=(1, 3),top_n=3,vectorizer=KeyphraseCountVectorizer())
    titles=[]
    for i in range(len(keywords)):
        title=keywords[i][0]
        titles.append(title)

    batch_dict = tokenizer(titles, padding=True, truncation=True, return_tensors='pt')
    outputs = model(**batch_dict)
    embeddings = average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
    embeddings = F.normalize(embeddings, p=2, dim=1)
    values = embeddings.tolist()

    catgories_list = []
    for value in values:
        try:
            response = index.query(vector=value,top_k=3,include_metadata=True)
        except:
            pinecone.init(api_key='b33ddf5d-5b1a-4d0e-9a3f-572008563791',environment='gcp-starter')
            index = pinecone.Index("wiki-index")
            response = index.query(vector=value,top_k=5,include_metadata=True)

        catgories = response['matches'][0]['metadata']['categories']
        catgories_list.append(catgories.split(','))
    flatList = [element for innerList in catgories_list for element in innerList]
    new_list = [item.replace("'", '') for item in flatList]
    a_list = [s.strip() for s in new_list]
    test_list = list(set(a_list))
    # result=[j for i in flatList for j in i]
    # res = [i for n, i in enumerate(result) if i not in result[:n]]
    return test_list

def main():
    st.title('Wiki Classifier')
    
    pnkc_no = st.text_input("Enter a pnkc number:")
    pat_text = st.text_area("Enter a text paragraph:")

    if st.button('Get Wiki categories'):

        if pnkc_no:
            text = get_pat_text(pnkc_no)
        else:
            text=pat_text

        st.write("Predicting Wiki Categories for text:",text[:200])
        start_time = time.time()
        titles,wiki_categories=get_wiki_category_aprch_1(text)
        end_time = time.time()
        st.write({f"Wiki_titles for {pnkc_no} Text":titles})
        st.write({f"Wiki_categories for {pnkc_no} Text":wiki_categories})
            
if __name__ == "__main__":
    main()