mb-rag 1.1.21__py3-none-any.whl → 1.1.22__py3-none-any.whl

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
+            device (str): Device to use. Default is cpu
             **kwargs: Additional arguments
         Returns:
             ChatHuggingFace: Chatbot model
@@ -225,22 +226,23 @@ class ModelFactory:
         from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
         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)
+        tokenizer = AutoTokenizer.from_pretrained(model_name,trust_remote_code=True)
         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,

mb_rag/version.py

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

{mb_rag-1.1.21.dist-info → mb_rag-1.1.22.dist-info}/METADATA

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

{mb_rag-1.1.21.dist-info → mb_rag-1.1.22.dist-info}/RECORD

@@ -1,7 +1,7 @@
 mb_rag/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-mb_rag/version.py,sha256=goWLkakwqod2v59-27LsKWvHHVJ1y6U_KJvRRJIJ56I,207
+mb_rag/version.py,sha256=MdYlquliW4HTyGehBHyYrWkbKgMYvz-ggotKwZvsUxE,207
 mb_rag/chatbot/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-mb_rag/chatbot/basic.py,sha256=gzxLZj96j_aZm3TXtryUBmGO90vNYGtg6TAQJsAmMz0,22432
+mb_rag/chatbot/basic.py,sha256=Bb3fAoYfkTPiB5vlQ9K4TfEO3m-TqukLhLzMbmdSOkc,22585
 mb_rag/chatbot/chains.py,sha256=vDbLX5R29sWN1pcFqJ5fyxJEgMCM81JAikunAEvMC9A,7223
 mb_rag/chatbot/prompts.py,sha256=n1PyiLbU-5fkslRv6aVOzt0dDlwya_cEdQ7kRnRhMuY,1749
 mb_rag/rag/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -9,7 +9,7 @@ mb_rag/rag/embeddings.py,sha256=KjBdekFDb5M3dRMco4r3dDMXMsG5dxdzKImuVIipsd0,2709
 mb_rag/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 mb_rag/utils/bounding_box.py,sha256=G0hdDam8QmYtD9lfwMeDHGm-TTo6KZg-yK5ESFL9zaM,8366
 mb_rag/utils/extra.py,sha256=spbFrGgdruNyYQ5PzgvpSIa6Nm0rn9bb4qc8W9g582o,2492
-mb_rag-1.1.21.dist-info/METADATA,sha256=A6YdCferOs7EC0kyDjDnEeO1grRwN7g1zZ_h9e3I0Zw,234
-mb_rag-1.1.21.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
-mb_rag-1.1.21.dist-info/top_level.txt,sha256=FIK1eAa5uYnurgXZquBG-s3PIy-HDTC5yJBW4lTH_pM,7
-mb_rag-1.1.21.dist-info/RECORD,,
+mb_rag-1.1.22.dist-info/METADATA,sha256=cghtqSh0h2irAQ3MZvHzRIiYwrPRyl5BPvMdzD34M2I,234
+mb_rag-1.1.22.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
+mb_rag-1.1.22.dist-info/top_level.txt,sha256=FIK1eAa5uYnurgXZquBG-s3PIy-HDTC5yJBW4lTH_pM,7
+mb_rag-1.1.22.dist-info/RECORD,,