Spaces:
Sleeping
Sleeping
File size: 3,352 Bytes
3206264 a61c2dd 3206264 fdc280f 3206264 35968cd 3206264 35968cd 3206264 35968cd 3206264 2a632a7 3206264 fdc280f 3206264 fdc280f 2a632a7 3206264 fdc280f 3206264 fdc280f 3206264 fdc280f 3206264 fdc280f 3206264 fdc280f 3206264 fdc280f 3206264 fdc280f 3206264 fdc280f 35968cd a61c2dd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
# aura-mind-glow/tools.py
import torch
from PIL import Image
from transformers import AutoProcessor
from unsloth import FastVisionModel
from langchain_community.vectorstores import FAISS
from vector_store import search_documents
from bigquery_search import search_bigquery_for_remedy
def create_plant_diagnosis_tool(model: FastVisionModel, processor: AutoProcessor):
"""Factory function to create the plant diagnosis tool."""
def diagnose_plant(file_path: str) -> str:
"""
Analyzes an image of a maize plant from a given file path and returns a diagnosis.
Args:
file_path: The local path to the image file.
"""
if model is None or processor is None or file_path is None:
return "Error: Vision model is not loaded or no file path was provided."
try:
image = Image.open(file_path)
except Exception as e:
return f"Error opening image file: {e}"
image = image.convert("RGB")
messages = [
{"role": "user", "content": [{"type": "text", "text": "What is the condition of this maize plant?"}, {"type": "image", "image": image}]}
]
text_prompt = processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = processor(text=text_prompt, images=image, return_tensors="pt").to(model.device)
with torch.inference_mode():
outputs = model.generate(**inputs, max_new_tokens=48, use_cache=True)
print(f"Model outputs: {outputs}")
response = processor.batch_decode(outputs, skip_special_tokens=True)[0]
# Clean up the model's output to get only the diagnosis
answer_start_index = response.rfind("model\n")
if answer_start_index != -1:
return response[answer_start_index + len("model\n"):
].strip()
return "Could not parse diagnosis from model output."
return diagnose_plant
def create_remedy_retrieval_tool(retriever: FAISS):
"""Factory function to create the remedy retrieval tool."""
def retrieve_remedy(diagnosis: str) -> str:
"""
Takes a plant health diagnosis (e.g., 'Phosphorus Deficiency') and returns a suggested remedy from a local knowledge base.
"""
if retriever is None:
return "Error: Knowledge base is not loaded."
retrieved_docs = retriever.invoke(diagnosis)
if retrieved_docs:
return retrieved_docs[0].page_content
return "No specific remedy found in the knowledge base for this condition."
return retrieve_remedy
def create_chroma_db_search_tool():
"""Factory function to create the ChromaDB search tool."""
def search_chroma_db(query: str) -> str:
"""
Searches the local ChromaDB vector store for a remedy based on a diagnosis query.
"""
results = search_documents(query)
return results[0] if results else "No remedy found in local knowledge base."
return search_chroma_db
def create_bigquery_search_tool():
"""Factory function to create the BigQuery search tool."""
def search_bigquery(query: str) -> str:
"""
Searches BigQuery for a remedy based on a diagnosis query.
"""
return search_bigquery_for_remedy(query)
return search_bigquery
|