Update app.py

app.py

@@ -1,47 +1,13 @@
+import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-example = """
-import sys
-import os
-import time
+def run(input):
+    return "Rofl: " + name
 
-sys.path.append(os.getcwd())
-
-from cluster.prepare_data import get_headers_pairs_list, write_dist_matrix
-from cluster.token_edit_distance import get_distance_matrix
-
-if len(sys.argv) < 3:
-    print(
-        "Too few arguments. You should provide: \n1. dataset_filename" +
-        "\n2. output_data_filename"
-    )
-    sys.exit()
-
-start = time.perf_counter()
-dataset_filename_ = sys.argv[1]
-output_data_filename_ = sys.argv[2]
-
-headers_pairs = get_headers_pairs_list(dataset_filename_, verbose=True)
-
-dist_matrix, max_dist = get_distance_matrix(
-    list(map(lambda x: x[1], headers_pairs)),
-    verbose=True
+app = gr.Interface(
+    fn=run,
+    inputs=["text"],
+    outputs=["text"]
 )
 
-write_dist_matrix(dist_matrix, max_dist, output_data_filename_, verbose=True)
-
-end = time.perf_counter()
-"""
-
-tokenizer = AutoTokenizer.from_pretrained('JetBrains/Mellum-4b-base')
-model = AutoModelForCausalLM.from_pretrained('JetBrains/Mellum-4b-base')
-encoded_input = tokenizer(example, return_tensors='pt', return_token_type_ids=False)
-input_len = len(encoded_input["input_ids"][0])
-out = model.generate(
-    **encoded_input,
-    max_new_tokens=100,
-)
-print("### Context")
-print(tokenizer.decode(out[0][:input_len]))
-print("### Prediction")
-print(tokenizer.decode(out[0][input_len:]))
+app.launch()