Spaces:
Sleeping
Sleeping
Carlos Salgado
commited on
Commit
•
d17ba2d
1
Parent(s):
9689d80
ignore flakes, add ingest
Browse files- .gitignore +2 -0
- generate_metadata.py +43 -34
- ingest.py +7 -0
.gitignore
CHANGED
@@ -5,3 +5,5 @@
|
|
5 |
.venv
|
6 |
.ipynb_checkpoints
|
7 |
|
|
|
|
|
|
5 |
.venv
|
6 |
.ipynb_checkpoints
|
7 |
|
8 |
+
flake.lock
|
9 |
+
flake.nix
|
generate_metadata.py
CHANGED
@@ -2,42 +2,51 @@ import os
|
|
2 |
import json
|
3 |
import openai
|
4 |
from dotenv import load_dotenv
|
|
|
5 |
|
6 |
from schema import Metadata, BimDiscipline
|
7 |
|
8 |
load_dotenv()
|
9 |
|
10 |
-
|
11 |
-
|
12 |
-
|
13 |
-
|
14 |
-
# Create client
|
15 |
-
client = openai.OpenAI(
|
16 |
-
|
17 |
-
|
18 |
-
)
|
19 |
-
|
20 |
-
# Call the LLM with the JSON schema
|
21 |
-
chat_completion = client.chat.completions.create(
|
22 |
-
|
23 |
-
|
24 |
-
|
25 |
-
|
26 |
-
|
27 |
-
|
28 |
-
|
29 |
-
|
30 |
-
|
31 |
-
|
32 |
-
|
33 |
-
]
|
34 |
-
)
|
35 |
-
|
36 |
-
created_user = json.loads(chat_completion.choices[0].message.content)
|
37 |
-
|
38 |
-
|
39 |
-
|
40 |
-
|
41 |
-
|
42 |
-
|
43 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
2 |
import json
|
3 |
import openai
|
4 |
from dotenv import load_dotenv
|
5 |
+
import argparse
|
6 |
|
7 |
from schema import Metadata, BimDiscipline
|
8 |
|
9 |
load_dotenv()
|
10 |
|
11 |
+
def extract_metadata(filename):
|
12 |
+
with open(filename, 'r') as f:
|
13 |
+
context = f.readlines()
|
14 |
+
|
15 |
+
# Create client
|
16 |
+
client = openai.OpenAI(
|
17 |
+
base_url="https://api.together.xyz/v1",
|
18 |
+
api_key=os.environ["TOGETHER_API_KEY"],
|
19 |
+
)
|
20 |
+
|
21 |
+
# Call the LLM with the JSON schema
|
22 |
+
chat_completion = client.chat.completions.create(
|
23 |
+
model="mistralai/Mixtral-8x7B-Instruct-v0.1",
|
24 |
+
response_format={"type": "json_object", "schema": Metadata.model_json_schema()},
|
25 |
+
messages=[
|
26 |
+
{
|
27 |
+
"role": "system",
|
28 |
+
"content": f"You are a helpful assistant that understands BIM documents and engineering disciplines. Your answer should be in JSON format and only include the title, a brief one-sentence summary, and the discipline the document belongs to, distinguishing between {[d.value for d in BimDiscipline]} based on the given document."
|
29 |
+
},
|
30 |
+
{
|
31 |
+
"role": "user",
|
32 |
+
"content": f"Analyze the provided document, which could be in either German or English. Extract the title, summarize it briefly in one sentence, and infer the discipline. Document:\n{' '.join(context)}"
|
33 |
+
}
|
34 |
+
]
|
35 |
+
)
|
36 |
+
|
37 |
+
created_user = json.loads(chat_completion.choices[0].message.content)
|
38 |
+
return created_user
|
39 |
+
|
40 |
+
if __name__ == "__main__":
|
41 |
+
parser = argparse.ArgumentParser(description="Generate metadata for a BIM document")
|
42 |
+
parser.add_argument("document", metavar="FILEPATH", type=str,
|
43 |
+
help="Path to the BIM document")
|
44 |
+
|
45 |
+
args = parser.parse_args()
|
46 |
+
|
47 |
+
if not os.path.exists(args.document) or not os.path.isfile(args.document):
|
48 |
+
print("File '{}' not found or not accessible.".format(args.document))
|
49 |
+
sys.exit(-1)
|
50 |
+
|
51 |
+
metadata = extract_metadata(args.document)
|
52 |
+
print(json.dumps(metadata, indent=2))
|
ingest.py
ADDED
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from langchain_community.document_loaders import UnstructuredPDFLoader
|
2 |
+
|
3 |
+
def ingest_pdf(path):
|
4 |
+
loader = UnstructuredPDFLoader()
|
5 |
+
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
|
6 |
+
|
7 |
+
return data
|