Update app.py

app.py

@@ -18,7 +18,6 @@ examples = [
   char *v11; // r13
   char *v12; // rax
   char v13[42]; // [rsp+Eh] [rbp-2Ah] BYREF
-
   v5 = (unsigned int)(a1 - 1);
   v6 = status;
   if ( (unsigned int)v5 <= 3 )
@@ -48,10 +47,9 @@ examples = [
 def predict_summary(code):
     global model
     global tokenizer
-    input = tokenizer('summarize: '+code,return_tensors='pt',max_length=max_input_length,truncation=True)
-    output = model.generate(**input,max_new_tokens=256)[0]
-    return tokenizer.decode(output,skip_special_tokens=True)
-
+    input = tokenizer('summarize: '+code, return_tensors='pt', max_length=max_input_length, truncation=True)
+    output = model.generate(**input, max_new_tokens=256)[0]
+    return tokenizer.decode(output, skip_special_tokens=True)
 
 # predict identifier (func name)
 def predict_identifier(code):
@@ -60,20 +58,35 @@ def predict_identifier(code):
     '''
     code should be like: "unsigned __int8 *__cdecl <func>(int *<var_0>,...){ return <func_1>(1);}"
     '''
-    input = tokenizer('identifier_predict: '+code,return_tensors='pt',max_length=max_input_length,truncation=True)
+    input = tokenizer('identifier_predict: '+code, return_tensors='pt', max_length=max_input_length, truncation=True)
     output = model.generate(**input)[0]
     return tokenizer.decode(output)
 
-# Create the Gradio interface
-iface = gr.Interface(
+# Create the Gradio interface for predicting identifiers
+identifier_iface = gr.Interface(
     fn=predict_identifier, 
     inputs="text", 
     outputs="text",
-    title="Predict identifiers",
-    description="Enter a prompt and see the model generate text.",
+    title="Predict Identifiers",
+    description="Enter a code snippet and see the model generate function identifiers.",
     examples=examples
 )
 
-# Launch the interface
-iface.launch()
-gr.load("models/ejschwartz/hext5").launch()
+# Create the Gradio interface for predicting summaries
+summary_iface = gr.Interface(
+    fn=predict_summary, 
+    inputs="text", 
+    outputs="text",
+    title="Predict Summary",
+    description="Enter a code snippet and see the model generate a summary.",
+    examples=examples
+)
+
+# Combine the interfaces into a single tabbed interface
+combined_iface = gr.TabbedInterface(
+    interface_list=[identifier_iface, summary_iface],
+    tab_names=["Predict Identifiers", "Predict Summary"]
+)
+
+# Launch the combined interface
+combined_iface.launch()