import os import requests import tellurium as te import tempfile import streamlit as st from langchain_text_splitters import CharacterTextSplitter from transformers import pipeline import chromadb # Constants and global variables GITHUB_OWNER = "sys-bio" GITHUB_REPO_CACHE = "BiomodelsCache" BIOMODELS_JSON_DB_PATH = "src/cached_biomodels.json" LOCAL_DOWNLOAD_DIR = tempfile.mkdtemp() cached_data = None db = None # Initialize Hugging Face model pipelines summarizer = pipeline("summarization", model="facebook/bart-large-cnn") llm = pipeline("text-generation", model="gpt2") def fetch_github_json(): url = f"https://api.github.com/repos/{GITHUB_OWNER}/{GITHUB_REPO_CACHE}/contents/{BIOMODELS_JSON_DB_PATH}" headers = {"Accept": "application/vnd.github+json"} response = requests.get(url, headers=headers) if response.status_code == 200: data = response.json() if "download_url" in data: file_url = data["download_url"] json_response = requests.get(file_url) return json_response.json() else: raise ValueError(f"Unable to fetch model DB from GitHub repository: {GITHUB_OWNER} - {GITHUB_REPO_CACHE}") else: raise ValueError(f"Unable to fetch model DB from GitHub repository: {GITHUB_OWNER} - {GITHUB_REPO_CACHE}") def search_models(search_str): global cached_data if cached_data is None: cached_data = fetch_github_json() query_text = search_str.strip().lower() models = {} for model_id, model_data in cached_data.items(): if 'name' in model_data: name = model_data['name'].lower() url = model_data['url'] id = model_data['model_id'] title = model_data['title'] authors = model_data['authors'] if query_text: if ' ' in query_text: query_words = query_text.split(" ") if all(word in ' '.join([str(v).lower() for v in model_data.values()]) for word in query_words): models[model_id] = { 'ID': model_id, 'name': name, 'url': url, 'id': id, 'title': title, 'authors': authors, } else: if query_text in ' '.join([str(v).lower() for v in model_data.values()]): models[model_id] = { 'ID': model_id, 'name': name, 'url': url, 'id': id, 'title': title, 'authors': authors, } return models def download_model_file(model_url, model_id): model_url = f"https://raw.githubusercontent.com/konankisa/BiomodelsStore/main/biomodels/{model_id}/{model_id}_url.xml" response = requests.get(model_url) if response.status_code == 200: os.makedirs(LOCAL_DOWNLOAD_DIR, exist_ok=True) file_path = os.path.join(LOCAL_DOWNLOAD_DIR, f"{model_id}.xml") with open(file_path, 'wb') as file: file.write(response.content) print(f"Model {model_id} downloaded successfully: {file_path}") return file_path else: raise ValueError(f"Failed to download the model from {model_url}") def convert_sbml_to_antimony(sbml_file_path, antimony_file_path): try: r = te.loadSBMLModel(sbml_file_path) antimony_str = r.getCurrentAntimony() with open(antimony_file_path, 'w') as file: file.write(antimony_str) print(f"Successfully converted SBML to Antimony: {antimony_file_path}") except Exception as e: print(f"Error converting SBML to Antimony: {e}") def split_biomodels(antimony_file_path): text_splitter = CharacterTextSplitter( separator=" // ", chunk_size=1000, chunk_overlap=20, length_function=len, is_separator_regex=False ) final_items = [] directory_path = os.path.dirname(os.path.abspath(antimony_file_path)) if not os.path.isdir(directory_path): print(f"Directory not found: {directory_path}") return final_items files = os.listdir(directory_path) for file in files: file_path = os.path.join(directory_path, file) try: with open(file_path, 'r') as f: file_content = f.read() items = text_splitter.create_documents([file_content]) for item in items: final_items.append(item) break except Exception as e: print(f"Error reading file {file_path}: {e}") return final_items def create_vector_db(final_items): global db client = chromadb.Client() db = client.create_collection( name="BioModelsRAG", metadata={"hnsw:space": "cosine"} ) documents = [] print("VectorDB successfully created.") for item in final_items: prompt = f""" Summarize the following segment of Antimony: {item} """ response = summarizer(prompt, max_length=150, min_length=30, do_sample=False) summary = response[0]['summary_text'] documents.append(summary) if final_items: db.add( documents=documents, ids=[f"id{i}" for i in range(len(final_items))] ) return db def generate_response(db, query_text, previous_context): query_results = db.query( query_texts=query_text, n_results=5, ) if not query_results.get('documents'): return "No results found." best_recommendation = query_results['documents'][0] prompt_template = f""" Using the context below, answer the following question: {query_text} Context: {previous_context} {best_recommendation} """ response = llm(prompt_template, max_length=150) final_response = response[0]['generated_text'] return final_response def streamlit_app(): st.title("BioModels Chat Interface") search_str = st.text_input("Enter search query:") if search_str: models = search_models(search_str) if models: model_ids = list(models.keys()) selected_models = st.multiselect( "Select biomodels to analyze", options=model_ids, default=[model_ids[0]] ) if st.button("Analyze Selected Models"): all_final_items = [] for model_id in selected_models: model_data = models[model_id] st.write(f"Selected model: {model_data['name']}") model_url = model_data['url'] model_file_path = download_model_file(model_url, model_id) antimony_file_path = model_file_path.replace(".xml", ".antimony") convert_sbml_to_antimony(model_file_path, antimony_file_path) final_items = split_biomodels(antimony_file_path) if not final_items: st.write("No content found in the biomodel.") continue all_final_items.extend(final_items) global db db = create_vector_db(all_final_items) if db: st.write("Models have been processed and added to the database.") user_query = st.text_input("Ask a question about the biomodels:") if user_query: if 'previous_context' not in st.session_state: st.session_state.previous_context = "" response = generate_response(db, user_query, st.session_state.previous_context) st.write(f"Response: {response}") st.session_state.previous_context += f"{response}\n" else: st.write("No models found for the given search query.") if __name__ == "__main__": streamlit_app()