Update app.py

app.py

@@ -1,29 +1,38 @@
 import gradio as gr
+import transformers
 
-# Load the pre-built interface
-interface = gr.load("models/alexue4/text-normalization-ru-new")
+# Load the model
+model_name = "alexue4/text-normalization-ru-new"
+model = transformers.AutoModelForSequenceClassification.from_pretrained(model_name)
+tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
 
-# Function to modify the behavior (example: limiting max tokens and setting examples)
+# Define the custom function
 def custom_function(input_text):
-    # Custom processing here (e.g., limit max tokens)
-    max_tokens = 250  # Set your desired max token limit
-    if len(input_text.split()) > max_tokens:
-        input_text = ' '.join(input_text.split()[:max_tokens])
+    max_tokens = 100  # Set your desired max token limit
+    tokens = tokenizer.tokenize(input_text)
     
-    # Call the original function from the interface
-    original_output = interface.predict(input_text)
-    return original_output
+    if len(tokens) > max_tokens:
+        tokens = tokens[:max_tokens]
+        input_text = tokenizer.convert_tokens_to_string(tokens)
+    
+    inputs = tokenizer(input_text, return_tensors="pt")
+    outputs = model(**inputs)
+    prediction = outputs.logits.argmax(dim=-1).item()
+    
+    return prediction
 
-# Customize the examples
 examples = [
     ["в 2006-2010 гг. Вася учился в МГУ"],
     ["я купил iphone 10X за 14990 руб без 3-x часов полдень и т.д."]
 ]
 
-# Define the new Gradio interface with custom function and examples
-gr.Interface(
+interface = gr.Interface(
     fn=custom_function,
     inputs="text",
     outputs="text",
     examples=examples
-).launch()
+)
+
+# Launch the interface
+interface.launch()
+