Update app.py

app.py

@@ -7,27 +7,18 @@ import gradio as gr
 import torch
 
 from transformers import (
+    SeamlessM4TTokenizer,
     SeamlessM4TProcessor,
     SeamlessM4TForTextToText,
-    SeamlessM4TTokenizer,           # <<< import the tokenizer class
     pipeline as hf_pipeline
 )
 
-# ————————————————————
-# 1) Load SeamlessM4T tokenizer (slow) and processor
-MODEL_NAME = "facebook/hf-seamless-m4t-medium"
-device     = "cuda" if torch.cuda.is_available() else "cpu"
-
-# load the slow tokenizer (no conversion attempted)
-tokenizer  = SeamlessM4TTokenizer.from_pretrained(MODEL_NAME, use_fast=False)
-
-# pass it into the processor so it won't try to convert
-processor  = SeamlessM4TProcessor.from_pretrained(
-    MODEL_NAME,
-    tokenizer=tokenizer
-)
-
-m4t_model  = SeamlessM4TForTextToText.from_pretrained(MODEL_NAME).to(device).eval()
+# 1) Load SeamlessM4T (slow tokenizer)
+MODEL = "facebook/hf-seamless-m4t-medium"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+tokenizer = SeamlessM4TTokenizer.from_pretrained(MODEL, use_fast=False)
+processor = SeamlessM4TProcessor.from_pretrained(MODEL, tokenizer=tokenizer)
+m4t_model = SeamlessM4TForTextToText.from_pretrained(MODEL).to(device).eval()
 
 def translate_m4t(text, src_iso3, tgt_iso3, auto_detect=False):
     src = None if auto_detect else src_iso3
@@ -35,120 +26,104 @@ def translate_m4t(text, src_iso3, tgt_iso3, auto_detect=False):
     tokens = m4t_model.generate(**inputs, tgt_lang=tgt_iso3)
     return processor.decode(tokens[0].tolist(), skip_special_tokens=True)
 
-# ————————————————————
-# 2) BERT‐based NER
-ner = hf_pipeline(
-    "ner",
-    model="dslim/bert-base-NER-uncased",
-    grouped_entities=True
-)
+# 2) NER pipeline
+ner = hf_pipeline("ner", model="dslim/bert-base-NER-uncased", grouped_entities=True)
 
-# ————————————————————
-# 3) Geocoding & POIs via OpenStreetMap
-def geocode(place: str):
-    resp = requests.get(
+# 3) Geocode & POIs
+def geocode(place):
+    r = requests.get(
         "https://nominatim.openstreetmap.org/search",
         params={"q": place, "format": "json", "limit": 1},
         headers={"User-Agent":"iVoiceContext/1.0"}
     ).json()
-    if not resp:
-        return None
-    return float(resp[0]["lat"]), float(resp[0]["lon"])
+    if not r: return None
+    return {"lat": float(r[0]["lat"]), "lon": float(r[0]["lon"])}
 
 def fetch_osm(lat, lon, osm_filter, limit=5):
-    query = f"""
-    [out:json][timeout:25];
-    (
-      node{osm_filter}(around:1000,{lat},{lon});
-      way{osm_filter}(around:1000,{lat},{lon});
-    );
-    out center {limit};
+    payload = f"""
+      [out:json][timeout:25];
+      ( node{osm_filter}(around:1000,{lat},{lon});
+        way{osm_filter}(around:1000,{lat},{lon}); );
+      out center {limit};
     """
-    r = requests.post("https://overpass-api.de/api/interpreter", data={"data": query})
-    elems = r.json().get("elements", [])
-    return [
-      {"name": e["tags"].get("name", "")}
-      for e in elems
-      if e.get("tags", {}).get("name")
-    ]
-
-# ————————————————————
+    resp = requests.post("https://overpass-api.de/api/interpreter", data={"data": payload})
+    elems = resp.json().get("elements", [])
+    return [{"name": e["tags"]["name"]} for e in elems if e.get("tags",{}).get("name")]
+
+# 4) Main function
 def get_context(text: str,
-                source_lang: str,  # always ISO639-3, e.g. "eng"
-                output_lang: str,  # always ISO639-3, e.g. "fra"
+                source_lang: str,  # ISO-639-3 e.g. "eng"
+                output_lang: str,  # ISO-639-3 e.g. "fra"
                 auto_detect: bool):
-    # 1) Ensure English text for NER
+
+    # a) Ensure English for NER
     if auto_detect or source_lang != "eng":
-        en_text = translate_m4t(text, source_lang, "eng", auto_detect=auto_detect)
+        en = translate_m4t(text, source_lang, "eng", auto_detect=auto_detect)
     else:
-        en_text = text
-
-    # 2) Run NER
-    ner_out = ner(en_text)
-    ents = { ent["word"]: ent["entity_group"] for ent in ner_out }
-
-    results = {}
-    for ent_text, label in ents.items():
-        if label == "LOC":
-            geo = geocode(ent_text)
+        en = text
+
+    # b) Extract unique entities
+    ner_out = ner(en)
+    seen, entities = set(), []
+    for ent in ner_out:
+        w, lbl = ent["word"], ent["entity_group"]
+        if w in seen: continue
+        seen.add(w)
+
+        if lbl == "LOC":
+            geo = geocode(w)
             if not geo:
-                results[ent_text] = {"type":"location","error":"could not geocode"}
+                obj = {"text": w, "label": lbl, "type": "location", "error": "could not geocode"}
             else:
-                lat, lon = geo
-                results[ent_text] = {
-                    "type":        "location",
-                    "restaurants": fetch_osm(lat, lon, '["amenity"="restaurant"]'),
-                    "attractions": fetch_osm(lat, lon, '["tourism"="attraction"]'),
+                obj = {
+                    "text": w,
+                    "label": lbl,
+                    "type": "location",
+                    "geo": geo,
+                    "restaurants": fetch_osm(geo["lat"], geo["lon"], '["amenity"="restaurant"]'),
+                    "attractions": fetch_osm(geo["lat"], geo["lon"], '["tourism"="attraction"]')
                 }
+
         else:
+            # PERSON/ORG/MISC → Wikipedia
             try:
-                summ = wikipedia.summary(ent_text, sentences=2)
-            except Exception:
+                summ = wikipedia.summary(w, sentences=2)
+            except:
                 summ = "No summary available."
-            results[ent_text] = {"type":"wiki","summary": summ}
+            obj = {"text": w, "label": lbl, "type": "wiki", "summary": summ}
 
-    if not results:
-        return {"error":"no entities found"}
+        entities.append(obj)
 
-    # 3) Translate all text fields → output_lang
+    # c) Translate all fields → output_lang
     if output_lang != "eng":
-        for info in results.values():
-            if info["type"] == "wiki":
-                info["summary"] = translate_m4t(
-                    info["summary"], "eng", output_lang, auto_detect=False
-                )
-            elif info["type"] == "location":
-                for key in ("restaurants","attractions"):
-                    info[key] = [
-                      {"name": translate_m4t(item["name"], "eng", output_lang)}
-                      for item in info[key]
+        for e in entities:
+            if e["type"] == "wiki":
+                e["summary"] = translate_m4t(e["summary"], "eng", output_lang)
+            elif e["type"] == "location":
+                for field in ("restaurants","attractions"):
+                    e[field] = [
+                        {"name": translate_m4t(item["name"], "eng", output_lang)}
+                        for item in e[field]
                     ]
 
-    return results
+    # d) Return only entities
+    return {"entities": entities}
 
-# ————————————————————
+# 5) Gradio interface
 iface = gr.Interface(
     fn=get_context,
     inputs=[
-      gr.Textbox(lines=3, placeholder="Enter text…"),
-      gr.Textbox(label="Source Language (ISO 639-3)"),
-      gr.Textbox(label="Target Language (ISO 639-3)"),
-      gr.Checkbox(label="Auto-detect source language")
+        gr.Textbox(lines=3, placeholder="Enter text…"),
+        gr.Textbox(label="Source Language (ISO 639-3)"),
+        gr.Textbox(label="Target Language (ISO 639-3)"),
+        gr.Checkbox(label="Auto-detect source language")
     ],
     outputs="json",
-    title="iVoice Translate + Context-Aware",
-    description=(
-      "1) Translate your text → English (if needed)\n"
-      "2) Extract LOC/PERSON/ORG via BERT-NER\n"
-      "3) Geocode LOC → fetch nearby restaurants & attractions\n"
-      "4) Fetch Wikipedia summaries\n"
-      "5) Translate **all** results → your target language"
-    )
+    title="iVoice Context-Aware",
+    description="Returns only the detected entities and their related info."
 ).queue()
 
 if __name__ == "__main__":
-    iface.launch(
-      server_name="0.0.0.0",
-      server_port=int(os.environ.get("PORT", 7860)),
-      share=True
-    )
+    iface.launch(server_name="0.0.0.0",
+                 server_port=int(os.environ.get("PORT", 7860)),
+                 share=True)