Upload 3 files

lib%2Fgraph_extract.py

@@ -0,0 +1,151 @@
+import json
+import re
+from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
+import torch
+import warnings
+import spaces
+
+flash_attn_installed = False
+# try:
+#     import subprocess
+#     print("Installing flash-attn...")
+#     subprocess.run(
+#     "pip install flash-attn --no-build-isolation",
+#     env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
+#     shell=True,
+#     )
+#     flash_attn_installed = True
+# except Exception as e:
+#     print(f"Error installing flash-attn: {e}")
+
+
+# Suppress specific warnings
+warnings.filterwarnings(
+    "ignore",
+    message="You have modified the pretrained model configuration to control generation.",
+)
+warnings.filterwarnings(
+    "ignore",
+    message="You are not running the flash-attention implementation, expect numerical differences.",
+)
+
+print("Initializing application...")
+
+model = AutoModelForCausalLM.from_pretrained(
+    "sciphi/triplex", 
+    trust_remote_code=True,
+    # attn_implementation="flash_attention_2" if flash_attn_installed else None,
+    torch_dtype=torch.bfloat16,
+    device_map="auto", 
+    low_cpu_mem_usage=True,#advised if any device map given
+).eval()
+
+tokenizer = AutoTokenizer.from_pretrained(
+    "sciphi/triplex", 
+    trust_remote_code=True,
+    # attn_implementation="flash_attention_2" if flash_attn_installed else None,
+        torch_dtype=torch.bfloat16,
+    )
+
+
+print("Model and tokenizer loaded successfully.")
+
+# Set up generation config
+generation_config = GenerationConfig.from_pretrained("sciphi/triplex")
+generation_config.max_length = 2048
+generation_config.pad_token_id = tokenizer.eos_token_id
+
+
+@spaces.GPU
+def triplextract(text, entity_types, predicates):
+    input_format = """Perform Named Entity Recognition (NER) and extract knowledge graph triplets from the text. NER identifies named entities of given entity types, and triple extraction identifies relationships between entities using specified predicates. Return the result as a JSON object with an "entities_and_triples" key containing an array of entities and triples.
+        
+        **Entity Types:**
+        {entity_types}
+        
+        **Predicates:**
+        {predicates}
+        
+        **Text:**
+        {text}
+        """
+    message = input_format.format(
+                entity_types = json.dumps({"entity_types": entity_types}),
+                predicates = json.dumps({"predicates": predicates}),
+                text = text)
+    
+    # message = input_format.format(
+    #     entity_types=entity_types, predicates=predicates, text=text
+    # )
+
+    messages = [{"role": "user", "content": message}]
+
+    print("Tokenizing input...")
+    input_ids = tokenizer.apply_chat_template(
+        messages, add_generation_prompt=True, return_tensors="pt"
+    ).to(model.device)
+
+    attention_mask = input_ids.ne(tokenizer.pad_token_id)
+
+    print("Generating output...")
+    try:
+        with torch.no_grad():
+            output = model.generate(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                generation_config=generation_config,
+                do_sample=True,
+            )
+
+        decoded_output = tokenizer.decode(output[0], skip_special_tokens=True)
+        print("Decoding output completed.")
+
+        return decoded_output
+    except torch.cuda.OutOfMemoryError as e:
+        print(f"CUDA out of memory error: {e}")
+        return "Error: CUDA out of memory."
+    except Exception as e:
+        print(f"Error in generation: {e}")
+        return f"Error in generation, please try again: {str(e)}"
+
+def parse_triples(prediction):
+    entities = {}
+    relationships = []
+
+    try:
+        data = json.loads(prediction)
+        items = data.get("entities_and_triples", [])
+    except json.JSONDecodeError:
+        json_match = re.search(r"```json\s*(.*?)\s*```", prediction, re.DOTALL)
+        if json_match:
+            try:
+                data = json.loads(json_match.group(1))
+                items = data.get("entities_and_triples", [])
+            except json.JSONDecodeError:
+                items = re.findall(r"\[(.*?)\]", prediction)
+        else:
+            items = re.findall(r"\[(.*?)\]", prediction)
+
+    for item in items:
+        if isinstance(item, str):
+            try:
+                if ":" in item:
+                    id, entity = item.split(",", 1)
+                    id = id.strip("[]").strip()
+                    entity_type, entity_value = entity.split(":", 1)
+                    entities[id] = {
+                        "type": entity_type.strip(),
+                        "value": entity_value.strip(),
+                    }
+                else:
+                    parts = item.split()
+                    if len(parts) >= 3:
+                        source = parts[0].strip("[]")
+                        relation = " ".join(parts[1:-1])
+                        target = parts[-1].strip("[]")
+                        relationships.append((source, relation.strip(), target))
+            except Exception as e:
+                # TODO: Handle gracefully
+                print(f"Error in processing: {item}: {e}")
+                
+    return entities, relationships

lib%2Fsamples.py

@@ -0,0 +1,46 @@
+from collections import namedtuple
+
+Snippet = namedtuple('Snippet', ['text_input', 'entity_types', 'predicates'])
+
+snippets = {
+    'paris': Snippet(
+        text_input="""Paris is the capital of France. It has a population of two million people. 
+        The Eiffel Tower, located in Paris, is a famous landmark with a height of 324 meters. 
+        Paris is known for its romantic atmosphere.""",
+        entity_types="LOCATION, POPULATION, STYLE",
+        predicates="HAS, IS"
+    ),
+
+    'dickens': Snippet(
+        text_input="""It was the best of times, it was the worst of times, it was the age of wisdom, 
+        it was the age of foolishness, it was the epoch of belief, it was the epoch of incredulity, 
+        it was the season of Light, it was the season of Darkness, it was the spring of hope, 
+        it was the winter of despair, we had everything before us, we had nothing before us, 
+        we were all going direct to Heaven, we were all going direct the other way – in short, 
+        the period was so far like the present period, that some of its noisiest authorities 
+        insisted on its being received, for good or for evil, in the superlative degree of comparison only.""",
+        entity_types="EMOTION, EVENT, OUTCOME, PLACE",
+        predicates="WAS, HAD, WERE, IS"
+    ),
+
+    'tech_company': Snippet(
+        text_input="""Apple Inc. was founded by Steve Jobs, Steve Wozniak, and Ronald Wayne in 1976. 
+        Headquartered in Cupertino, California, Apple designs and produces consumer electronics, 
+        software, and online services. The company's flagship products include the iPhone smartphone, 
+        iPad tablet, and Mac personal computer. As of 2023, Apple has over 150000 employees worldwide 
+        and generates annual revenue exceeding $350 billion.""",
+        entity_types="COMPANY, PERSON, PRODUCT, LOCATION, DATE, EVENT",
+        predicates="FOUNDED, PRODUCES, HAS, IN, EMPLOYEES, "
+    ),
+
+    'climate_change': Snippet(
+        text_input="""Global warming is causing significant changes to Earth's climate. The average global 
+        temperature has increased by approximately 1C since the pre-industrial era. This warming is 
+        primarily caused by human activities, particularly the emission of greenhouse gases like carbon dioxide. 
+        The Paris Agreement, signed in 2015, aims to limit global temperature increase to well below 2°C above 
+        pre-industrial levels. To achieve this goal, many countries are implementing policies to reduce carbon 
+        emissions and transition to renewable energy sources.""",
+        entity_types="PHENOMENON, PLANET, TEMPERATURE, CAUSE, CHEMICAL, AGREEMENT, DATE, GOAL, POLICY",
+        predicates="CAUSES, INCREASED_BY, CAUSED_BY, SIGNED_IN, AIMS_TO, IMPLEMENTING"
+    )
+}