You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
ClinicalTrialsDataProcessing/llm-extractor/llm_extractor.py

56 lines
1.2 KiB
Python

import ollama
import psycopg
from psycopg.rows import dict_row
from typing import List, Dict
def fetch_all_rows(conn_params: dict) -> List[Dict]:
# Establish a connection to the PostgreSQL database
conn = psycopg.connect(**conn_params, row_factory=dict_row)
cursor = conn.cursor()
# Define your SQL query to select all rows from the table
sql_query = "SELECT * FROM public.primary_design_outcomes;"
# Execute the query
cursor.execute(sql_query)
# Fetch all rows from the result set
rows = cursor.fetchall()
# Close the cursor and connection
cursor.close()
conn.close()
return rows
# Example usage
conn_params = {
"dbname": "aact_db",
"user": "root",
"password": "root",
"host": "localhost",
"port": "5432"
}
outcome_description = '''
Measure: {measure}
Time Frame: {time_frame}
Description: {description}
'''
if __name__ == "__main__":
#check for model
#get information
rows_dicts = fetch_all_rows(conn_params)
for row in rows_dicts[:3]:
text_data = outcome_description.format(**row)
r = ollama.generate(model='youainti/llama3.1-extractor:2024-08-28.2', prompt=text_data)
print(text_data)
print(r["response"])