#!/usr/bin/env python # coding: utf-8 # In[10]: import pandas as pd import os import torch from transformers import T5Tokenizer, T5ForConditionalGeneration from transformers.optimization import Adafactor import time import warnings import random warnings.filterwarnings('ignore') import re def strip_html(text): return re.sub('<[^<]+?>', '', text) # In[5]: train_columns = ['round_amount', 'round_date', 'stage', 'investee', 'investee_description', 'investee_country', 'investee_region', 'investee_subregion', 'investee_vertical', 'investee_industry', 'investor_list', 'previous_investors', 'prior_funding'] train = pd.read_csv("train.csv") # In[6]: train.publication_timestamp = pd.to_datetime(train.publication_timestamp) # In[7]: input_text = train[train_columns].to_dict(orient='records') train_df = train[['title']].rename(columns={'title':'target_text'}) train_df['input_text'] = input_text train_df['prefix'] = 'tia' train_df.input_text = train_df.input_text.astype(str) # In[8]: if torch.cuda.is_available(): dev = torch.device("cuda:0") print("Running on the GPU") else: dev = torch.device("cpu") print("Running on the CPU") # In[ ]: tokenizer = T5Tokenizer.from_pretrained('google/t5-v1_1-base') model = T5ForConditionalGeneration.from_pretrained('t5-v1_1-base_tia/', local_files_only=True) #moving the model to device(GPU/CPU) model.to(dev) # In[12]: vi_table = train[['investee_industry', 'investee_vertical']].drop_duplicates() # In[13]: def update_industry(value): verticals = list(vi_table[vi_table['investee_industry'] == value]['investee_vertical'].values) return verticals[0] def update_vertical(value): industries = list(vi_table[vi_table['investee_vertical'] == value]['investee_industry'].values) return industries[0] # In[ ]: update_industry('Green') # In[ ]: update_vertical('Clean tech') # In[ ]: import gradio as gr # In[ ]: num_return_sequences = 5 # In[ ]: def generate_headline(stage, investee_country, investee_subregion, investee_region, investee_vertical, investee_industry, round_amount, investee, investee_description, investor_list, previous_investors, other_values): full_df = other_values.set_index("key").T full_df['stage'] = stage full_df['investee_country'] = investee_country full_df['investee_subregion'] = investee_subregion full_df['investee_region'] = investee_region full_df['investee_vertical'] = investee_vertical full_df['investee_industry'] = investee_industry full_df['round_amount'] = str(float(round_amount)) full_df['investee'] = investee full_df['investee_description'] = investee_description full_df['investor_list'] = investor_list full_df['previous_investors'] = previous_investors random_set =full_df[['round_amount', 'round_date', 'stage', 'investee', 'investee_description', 'investee_country', 'investee_region', 'investee_subregion', 'investee_vertical', 'investee_industry', 'investor_list', 'previous_investors', 'prior_funding']].to_json(orient="records") # print(random_set) input_ids = tokenizer.encode(f"tia: {{{random_set}}}", return_tensors="pt") # Batch size 1 input_ids=input_ids.to(dev) outputs = model.generate(input_ids) # text_output = tokenizer.decode(outputs[0]) # Single output text_outputs = model.generate(inputs=input_ids, do_sample=True, num_beams=2, num_return_sequences=num_return_sequences, repetition_penalty=5.0) outputs = [strip_html(tokenizer.decode(o)) for o in text_outputs] return "\n".join(outputs) # In[ ]: other_columns = ['round_date', 'prior_funding'] # In[ ]: train.sample(1)[other_columns].T.reset_index().values # In[ ]: print(train.query("investee == 'NOSH'")['title'].head(1).T) train.query("investee == 'NOSH'")[train_columns].head(1).T # In[ ]: fake_data = { "round_amount":1000000.0, "round_date":"2018-09-26", "stage":"Pre-series A", "investee":"NOSH", "investee_description":"NOSH makes and delivers ready-to-eat meals in Hong Kong.", "investee_country":"Hong Kong", "investee_region":"Asia", "investee_subregion":"Eastern Asia", "investee_vertical":"Food tech", "investee_industry":"Restaurants & Food", "investor_list":["Alibaba Entrepreneurs Fund (阿里巴巴创业者基金)"], "previous_investors":"", "prior_funding":1000000.0 } # In[ ]: pd.DataFrame([fake_data]).T # In[ ]: demo = gr.Blocks() random_sample = train[train_columns].sample(1) random_sample = pd.DataFrame([fake_data]) stage = gr.Dropdown(label="stage", choices=list(train[train_columns].stage.unique())) investee_country = gr.Dropdown(label="investee_country", choices=list(train[train_columns].investee_country.unique()), value=random_sample.investee_country.values[0]) investee_subregion = gr.Dropdown(label="investee_subregion", choices=list(train[train_columns].investee_subregion.unique()), value=random_sample.investee_subregion.values[0]) investee_region = gr.Dropdown(label="investee_region", choices=list(train[train_columns].investee_region.unique()), value=random_sample.investee_region.values[0]) investee_vertical = gr.Dropdown(label="investee_vertical", choices=list(train[train_columns].investee_vertical.unique()), value=random_sample.investee_vertical.values[0]) investee_industry = gr.Dropdown(label="investee_industry", choices=list(train[train_columns].investee_industry.unique()), value=random_sample.investee_industry.values[0]) if pd.isnull(random_sample.round_amount.values[0]): rand_amount = 0 else: rand_amount = random_sample.round_amount.values[0] round_amount = gr.Slider(label="round_amount", minimum=100000, maximum=200000000, value=rand_amount, step=100000) investee = gr.Textbox(label="investee", value=random_sample.investee.values[0]) investee_description = gr.Textbox(label="investee_description", value=random_sample.investee_description.values[0]) investor_list = gr.Textbox(label="investor_list", value=random_sample.investor_list.values[0]) previous_investors = gr.Textbox(label="previous_investors", value=random_sample.previous_investors.values[0]) other_values = gr.Dataframe( headers=['key', 'value'], value=[['round_date', random_sample.round_date.values[0]], ['prior_funding', random_sample.prior_funding.values[0]]] ) out = gr.Textbox(max_lines=num_return_sequences) with demo: gr.Markdown("Enter funding data to generate news headline.") inputs=[stage, investee_country, investee_subregion, investee_region, investee_vertical, investee_industry, round_amount, investee, investee_description, investor_list, previous_investors, other_values] investee_industry.change(fn=update_industry, inputs=investee_industry, outputs=investee_vertical) investee_vertical.change(fn=update_vertical, inputs=investee_vertical, outputs=investee_industry) gr.Interface(fn=generate_headline, inputs=inputs, outputs=out, live=True) description="Enter funding data to generate news headline.", live=True demo.launch()