Update app.py

app.py

@@ -1,38 +1,59 @@
 import gradio as gr
 import onnxruntime as ort
-from transformers import AutoTokenizer, pipeline
+from transformers import AutoTokenizer
 import numpy as np
 
-MODEL_FILE = "./model.onnx"
-session = ort.InferenceSession(MODEL_FILE)
+# Initialize models
+context_model_file = "./bart-base-nmli.onnx"
+translation_model_file = "./model.onnx"
+
+# Create inference sessions for both models
+context_session = ort.InferenceSession(context_model_file)
+translation_session = ort.InferenceSession(translation_model_file)
+
+# Load tokenizer for translation model
 tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-fr")
 
 labels = [
-    "general", "pharma", "legal", "technical", "UI", "user interface", "medicine", "information technology", "IT", "marketing", 
-    "e-commerce", "programming", "website", "html", "keywords", "food commerce", 
-    "personal development", "literature", "poetry", "physics", "chemistry", "biology", 
-    "theater", "finance", "sports", "education", "politics", "economics", "art", 
-    "history", "music", "gaming", "aerospace", "engineering", "robotics", "travel", 
-    "tourism", "healthcare", "psychology", "environment", "fashion", "design", 
-    "real estate", "retail", "news", "entertainment", "social media", "automotive", 
-    "machine learning", "cryptocurrency", "blockchain", "philosophy", "anthropology", 
-    "archaeology", "data science"
+    "general", "pharma", "legal", "technical", "UI", "user interface", "medicine", 
+    "information technology", "IT", "marketing", "e-commerce", "programming", 
+    "website", "html", "keywords", "food commerce", "personal development", 
+    "literature", "poetry", "physics", "chemistry", "biology", "theater", 
+    "finance", "sports", "education", "politics", "economics", "art", "history",
+    "music", "gaming", "aerospace", "engineering", "robotics", "travel", "tourism",
+    "healthcare", "psychology", "environment", "fashion", "design", "real estate",
+    "retail", "news", "entertainment", "social media","automotive", "machine learning",
+    "cryptocurrency","blockchain","philosophy","anthropology","archaeology","data science"
 ]
 
-#context_pipeline = pipeline("zero-shot-classification", model="MoritzLaurer/deberta-v3-base-mnli-fever-anli")
-context_pipeline = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
-
 def detect_context(input_text):
-    result = context_pipeline(input_text, candidate_labels=labels)
-    contexts = [label for label, score in zip(result["labels"], result["scores"]) if label in labels and score > 0.1]
-    print(contexts)
-    return contexts or ["general"]
+    # Tokenize input text
+    inputs = tokenizer(input_text, return_tensors="np")
+    
+    # Prepare input for ONNX model
+    input_ids = inputs["input_ids"].astype(np.int64)
+    
+    # Run inference with context model
+    outputs = context_session.run(None, {"input_ids": input_ids})
+    
+    # Assuming the output is logits for each label
+    scores = outputs[0]
+    
+    # Get the top label based on scores
+    top_label_index = np.argmax(scores, axis=1)[0]
+    
+    # Map index to label
+    detected_context = labels[top_label_index]
+    
+    print(detected_context)
+    return detected_context
 
-
-def translate_text(input_text, context):
+def translate_text(input_text):
     tokenized_input = tokenizer(
-        input_text, return_tensors="np", padding=True, truncation=True, max_length=512
+        input_text, return_tensors="np",
+        padding=True, truncation=True, max_length=512
     )
+    
     input_ids = tokenized_input["input_ids"].astype(np.int64)
     attention_mask = tokenized_input["attention_mask"].astype(np.int64)
 
@@ -40,7 +61,7 @@ def translate_text(input_text, context):
     decoder_input_ids = np.array([[decoder_start_token_id]], dtype=np.int64)
 
     for _ in range(512):
-        outputs = session.run(
+        outputs = translation_session.run(
             None,
             {
                 "input_ids": input_ids,
@@ -60,16 +81,14 @@ def translate_text(input_text, context):
 
     return tokenizer.decode(decoder_input_ids[0], skip_special_tokens=True)
 
-
 def process_request(input_text):
-    contexts = detect_context(input_text)
-    translation = translate_text(input_text, contexts[0])  # Use the first detected context
+    context = detect_context(input_text)
+    translation = translate_text(input_text)  # Translate without needing to pass context explicitly
     return translation
 
-
 gr.Interface(
     fn=process_request,
     inputs="text",
     outputs="text",
     live=True
-).launch()
+).launch()