|
import os |
|
import streamlit as st |
|
|
|
from datasets import load_dataset |
|
import re |
|
import time |
|
import requests |
|
from bs4 import BeautifulSoup |
|
import pandas as pd |
|
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') |
|
|
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained('intfloat/e5-base') |
|
model = AutoModel.from_pretrained('intfloat/e5-base') |
|
|
|
|
|
|
|
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') |
|
|
|
|
|
categories_section = soup.find("div", {"class": "mw-normal-catlinks"}) |
|
|
|
if categories_section: |
|
|
|
categories = [cat.text for cat in categories_section.find("ul").find_all("li")] |
|
|
|
|
|
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 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): |
|
|
|
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)) |
|
|
|
|
|
return test_list |
|
|
|
def main(): |
|
st.title('wikipedia-titles-category-generator') |
|
|
|
text = st.text_area("Enter a text paragraph:") |
|
|
|
if st.button('Get Wiki categories'): |
|
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 Text":titles}) |
|
st.write({f"Wiki_categories for Text":wiki_categories}) |
|
|
|
if __name__ == "__main__": |
|
main() |
|
|
