File size: 4,239 Bytes
7ef595f 1094480 7ef595f 757c949 7ef595f 757c949 7ef595f f5a3078 7ef595f 757c949 7ef595f 315a464 7ef595f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 |
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')
# 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
# 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('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()
|