Update app.py

app.py

@@ -24,10 +24,9 @@ context_pipeline = pipeline("zero-shot-classification", model="facebook/bart-lar
 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 score > 0.1]
-    print(contexts)
-    return contexts if contexts else ["general"]
+    return contexts[0] if contexts else "general"
 
-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
     )
@@ -37,35 +36,4 @@ def translate_text(input_text, context):
     decoder_start_token_id = tokenizer.cls_token_id or tokenizer.pad_token_id
     decoder_input_ids = np.array([[decoder_start_token_id]], dtype=np.int64)
 
-    for _ in range(512):
-        outputs = session.run(
-            None,
-            {
-                "input_ids": input_ids,
-                "attention_mask": attention_mask,
-                "decoder_input_ids": decoder_input_ids,
-            }
-        )
-
-        logits = outputs[0]
-        next_token_id = np.argmax(logits[:, -1, :], axis=-1).item()
-        decoder_input_ids = np.concatenate(
-            [decoder_input_ids, np.array([[next_token_id]], dtype=np.int64)], axis=1
-        )
-
-        if next_token_id == tokenizer.eos_token_id:
-            break
-
-    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
-    return translation, contexts[0]
-
-gr.Interface(
-    fn=process_request,
-    inputs="text",
-    outputs="text",
-    live=True
 ).launch()