mb-rag 1.1.21__tar.gz → 1.1.23__tar.gz

{mb_rag-1.1.21 → mb_rag-1.1.23}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: mb_rag
-Version: 1.1.21
+Version: 1.1.23
 Summary: RAG function file
 Author: ['Malav Bateriwala']
 Requires-Python: >=3.8

{mb_rag-1.1.21 → mb_rag-1.1.23}/mb_rag/chatbot/basic.py

@@ -204,12 +204,13 @@ class ModelFactory:
 
     @classmethod
     def create_hugging_face(cls, model_name: str = "Qwen/Qwen2.5-VL-7B-Instruct",model_function: str = "image-text-to-text",
-                            **kwargs) -> Any:
+                            device='cpu',**kwargs) -> Any:
         """
         Create and load hugging face model.
         Args:
             model_name (str): Name of the model
-            model_function (str): model function
+            model_function (str): model function. Default is image-text-to-text.
+            device (str): Device to use. Default is cpu
             **kwargs: Additional arguments
         Returns:
             ChatHuggingFace: Chatbot model
@@ -222,30 +223,38 @@ class ModelFactory:
             raise ImportError("Torch package not found. Please install it using: pip install torch")
 
         from langchain_huggingface import HuggingFacePipeline
-        from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+        from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, AutoModelForImageTextToText
         import torch
         
-        device = "cuda" if torch.cuda.is_available() else "cpu"
+        device = torch.device(device) if torch.cuda.is_available() else torch.device("cpu")
         
         temperature = kwargs.pop("temperature", 0.7)
         max_length = kwargs.pop("max_length", 1024)
         
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-            device_map=device,
-            **kwargs
-        )
+        tokenizer = AutoTokenizer.from_pretrained(model_name,trust_remote_code=True)
+        if model_function == "image-text-to-text":
+            model = AutoModelForImageTextToText.from_pretrained(
+                model_name,
+                torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+                device_map=device,
+                trust_remote_code=True,
+                **kwargs
+            )
+        else:
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+                device_map=device,
+                trust_remote_code=True,
+                **kwargs)
         
         # Create pipeline
         pipe = pipeline(
-            "text-generation",
+            model_function,
             model=model,
             tokenizer=tokenizer,
             max_length=max_length,
-            temperature=temperature,
-            device=device
+            temperature=temperature
         )
         
         # Create and return LangChain HuggingFacePipeline

{mb_rag-1.1.21 → mb_rag-1.1.23}/mb_rag/version.py

@@ -1,5 +1,5 @@
 MAJOR_VERSION = 1
 MINOR_VERSION = 1
-PATCH_VERSION = 21
+PATCH_VERSION = 23
 version = '{}.{}.{}'.format(MAJOR_VERSION, MINOR_VERSION, PATCH_VERSION)
 __all__  = ['MAJOR_VERSION', 'MINOR_VERSION', 'PATCH_VERSION', 'version']

{mb_rag-1.1.21 → mb_rag-1.1.23}/mb_rag.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: mb_rag
-Version: 1.1.21
+Version: 1.1.23
 Summary: RAG function file
 Author: ['Malav Bateriwala']
 Requires-Python: >=3.8