Update app.py

app.py

@@ -1,11 +1,10 @@
 import os
 import warnings
 from typing import *
-from dotenv import load_dotenv
-from transformers import logging
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from langchain.llms import HuggingFacePipeline
 
 from langgraph.checkpoint.memory import MemorySaver
-from langchain_openai import ChatOpenAI
 
 from interface import create_demo
 from medrax.agent import *
@@ -13,21 +12,6 @@ from medrax.tools import *
 from medrax.utils import *
 
 warnings.filterwarnings("ignore")
-logging.set_verbosity_error()
-load_dotenv() 
-
-# Set environment variables explicitly to ensure they're available
-api_key = os.getenv("OPENAI_API_KEY")
-base_url = os.getenv("OPENAI_BASE_URL")
-
-if not api_key:
-    raise ValueError("OPENAI_API_KEY not found in environment variables")
-if not base_url:
-    raise ValueError("OPENAI_BASE_URL not found in environment variables")
-
-# Set them in environment for libraries that might read directly from os.environ
-os.environ["OPENAI_API_KEY"] = api_key
-os.environ["OPENAI_BASE_URL"] = base_url
 
 def initialize_agent(
     prompt_file,
@@ -35,7 +19,6 @@ def initialize_agent(
     model_dir="./model-weights",
     temp_dir="temp",
     device="cuda",
-    model="qwen/qwen2.5-vl-3b-instruct:free",
     temperature=0.7,
     top_p=0.95
 ):
@@ -69,16 +52,24 @@ def initialize_agent(
             tools_dict[tool_name] = all_tools[tool_name]()
 
     checkpointer = MemorySaver()
-    
-    # Explicitly pass the API key and base URL
-    model = ChatOpenAI(
-        model_name=model,
-        api_key=api_key,
-        base_url=base_url,
+
+    # Load local Hugging Face model
+    hf_model_id = "mistralai/Mistral-7B-Instruct-v0.2"
+    tokenizer = AutoTokenizer.from_pretrained(hf_model_id)
+    raw_model = AutoModelForCausalLM.from_pretrained(hf_model_id, device_map="auto")
+
+    pipe = pipeline(
+        "text-generation",
+        model=raw_model,
+        tokenizer=tokenizer,
+        max_new_tokens=512,
         temperature=temperature,
         top_p=top_p,
+        return_full_text=False,
     )
 
+    model = HuggingFacePipeline(pipeline=pipe)
+
     agent = Agent(
         model,
         tools=list(tools_dict.values()),
@@ -113,12 +104,9 @@ if __name__ == "__main__":
         model_dir="./model-weights",
         temp_dir="temp",
         device="cuda",
-        model="qwen/qwen2.5-vl-3b-instruct:free",
         temperature=0.7,
         top_p=0.95
     )
 
     demo = create_demo(agent, tools_dict)
-    # demo.launch(server_name="0.0.0.0", server_port=8585, share=True)
-    # demo.launch(debug=True, queue=True, ssr_mode=False)
-    demo.launch(debug=True, ssr_mode=False)
+    demo.launch(debug=True, ssr_mode=False)