Update predefined.py

predefined.py

@@ -475,8 +475,9 @@ def ldr_new(I,B,M,N_df,WT,batch_size=32):
     return(dt_ldr)
 
 
-
-def get_output_new(w,wt,I_b=n_Identities,B_b=n_Behaviors,M_b=n_Modifiers,batch_sz=300,batch_num=1):
+cols=['EEMA', 'EPMA', 'EAMA', 'EEA', 'EPA', 'EAA', 'EEB', 'EPB', 'EAB',
+       'EEMO', 'EPMO', 'EAMO', 'EEO', 'EPO', 'EAO',  'ModA', 'Actor', 'Behavior', 'ModO', 'Object']
+def get_output_new(w,wt,I_b=n_Identities,B_b=n_Behaviors,M_b=n_Modifiers,batch_sz=300,batch_num=1,columnss=cols,cus_col=1):
     
     df=pd.DataFrame()
     for i in range(batch_num):
@@ -486,17 +487,22 @@ def get_output_new(w,wt,I_b=n_Identities,B_b=n_Behaviors,M_b=n_Modifiers,batch_s
         if wt=='identity':
             df_identity=pd.concat([Identities,new_df],axis=0)
             df2=out_df(data=q,predictions=preds,df_ident=df_identity)
+            if cus_col:
+                cols=[ 'EEA', 'EPA', 'EAA', 'ModA', 'Actor', 'Behavior', 'ModO', 'Object']
+            
 
         if wt=='behavior': 
             df_behavior=pd.concat([Behaviors,new_df],axis=0)
             df2=out_df(data=q,predictions=preds,df_beh=df_behavior)
-            
+            if cus_col:
+                cols=['EEB', 'EPB', 'EAB',  'ModA', 'Actor', 'Behavior', 'ModO', 'Object']            
         if wt=='modifier': 
             df_modifier=pd.concat([Modifiers,new_df],axis=0)
             df2=out_df(data=q,predictions=preds,df_mod=df_modifier)
-            
+            if cus_col:
+                cols=['EEMA', 'EPMA', 'EAMA',  'ModA', 'Actor', 'Behavior', 'ModO', 'Object']
         df=pd.concat([df,df2],axis=0)
-    return(df)
+    return(df[columnss)
 
 def gen_new(Identity,Behavior,Modifier,n_df,word_type):
         if word_type=='identity':
@@ -614,5 +620,3 @@ def freeze_header(df, num_rows=30, num_columns=10, step_rows=1,
         
 
 
-cols=['EEMA', 'EPMA', 'EAMA', 'EEA', 'EPA', 'EAA', 'EEB', 'EPB', 'EAB',
-       'EEMO', 'EPMO', 'EAMO', 'EEO', 'EPO', 'EAO',  'ModA', 'Actor', 'Behavior', 'ModO', 'Object']