Update app.py

app.py

@@ -1,20 +1,17 @@
 import json
 import numpy as np
 import gradio as gr
-import os
 from sentence_transformers import SentenceTransformer
-from huggingface_hub import InferenceClient
+from transformers import T5Tokenizer, T5ForConditionalGeneration
 
-# Read the 'token' secret and set it as HF_TOKEN
-os.environ["HF_TOKEN"] = os.getenv("token", "Not set")
-
-# Debug: Confirm HF_TOKEN is set
-print("HF_TOKEN:", os.getenv("HF_TOKEN", "Not set"))
-
-# 1) Load the embedding model
+# 1) Load the embedding model for retrieval
 embedder = SentenceTransformer('all-MiniLM-L6-v2')
 
-# 2) Load JSON knowledge bases
+# 2) Load the T5 model and tokenizer for translation
+tokenizer = T5Tokenizer.from_pretrained('t5-small')
+model = T5ForConditionalGeneration.from_pretrained('t5-small')
+
+# 3) Load JSON knowledge bases
 def load_docs(path, fmt):
     with open(path, 'r') as f:
         data = json.load(f)
@@ -30,16 +27,13 @@ truck_docs = load_docs(
 )
 docs = call_docs + truck_docs
 
-# 3) Pre-compute embeddings for all documents
+# 4) Pre-compute embeddings for all documents
 doc_embeddings = embedder.encode(
     docs,
     convert_to_numpy=True,
     normalize_embeddings=True
 )
 
-# 4) Initialize the Hugging Face Inference client with t5-small
-client = InferenceClient(model="t5-small")
-
 def rag_translate(query, top_k):
     # 4.1 Retrieval step
     q_emb = embedder.encode(
@@ -52,23 +46,31 @@ def rag_translate(query, top_k):
     retrieved = [docs[i] for i in idxs]
     retrieved_text = "\n".join(retrieved)
 
-    # 4.2 Translation step with error handling
+    # 4.2 Translation step using local t5-small
     try:
+        # Prepare input for translation
+        spanish_input = f"translate English to Spanish: {retrieved_text}"
+        french_input = f"translate English to French: {retrieved_text}"
+
         # Translate to Spanish
-        spanish_result = client.translation(
-            retrieved_text,
-            src_lang="en",
-            tgt_lang="es"
+        spanish_inputs = tokenizer(spanish_input, return_tensors="pt", padding=True, truncation=True, max_length=512)
+        spanish_outputs = model.generate(
+            spanish_inputs["input_ids"],
+            max_length=512,
+            num_beams=4,
+            early_stopping=True
         )
-        spanish_text = spanish_result["translation_text"] if isinstance(spanish_result, dict) else spanish_result
+        spanish_text = tokenizer.decode(spanish_outputs[0], skip_special_tokens=True)
 
         # Translate to French
-        french_result = client.translation(
-            retrieved_text,
-            src_lang="en",
-            tgt_lang="fr"
+        french_inputs = tokenizer(french_input, return_tensors="pt", padding=True, truncation=True, max_length=512)
+        french_outputs = model.generate(
+            french_inputs["input_ids"],
+            max_length=512,
+            num_beams=4,
+            early_stopping=True
         )
-        french_text = french_result["translation_text"] if isinstance(french_result, dict) else french_result
+        french_text = tokenizer.decode(french_outputs[0], skip_special_tokens=True)
 
         # Combine translations
         translations = f"Spanish:\n{spanish_text}\n\nFrench:\n{french_text}"
@@ -107,7 +109,7 @@ with gr.Blocks() as demo:
         """
         **How it works**  
         1) Retrieval: fetch top-k snippets from JSON KBs.  
-        2) Translation: translate retrieved content to Spanish and French using t5-small.  
+        2) Translation: translate retrieved content to Spanish and French using t5-small locally.  
         """
     )