phi2 bug fix

README.md

@@ -4,6 +4,8 @@
 
 
 [![Downloads](https://static.pepy.tech/badge/nl2query)](https://pepy.tech/project/nl2query)
+[![Build Status][build-image]][build-url]
+[![][stars-image]][stars-url]
 [![PyPI version][pypi-image]][pypi-url]
 [![Support Python versions][versions-image]][versions-url]
 [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
@@ -43,6 +45,8 @@ Suppose you want to convert the textual question to Mongo query, follow the inst
 
 The generate_query method takes a textual query and returns a MongoDB query. It also accepts optional parameters to control the generation process, such as num_beams, max_length, repetition_penalty, length_penalty, early_stopping, top_p, top_k, and num_return_sequences.
 
+NOTE: GPU will be required to run Phi2 as quantization is enabled using *load_in_4bit*.
+
 ```py
 from nl2query import MongoQuery
 import pymongo # import if performing analysis using python client
@@ -202,3 +206,7 @@ Refer to the [CHANGELOG.md](CHANGELOG.md) file.
 [pypi-url]: https://pypi.org/project/nl2query/
 [versions-image]: https://img.shields.io/pypi/pyversions/nl2query
 [versions-url]: https://pypi.org/project/nl2query/
+[build-image]: https://github.com/Chirayu-Tripathi/nl2query/actions/workflows/build.yaml/badge.svg
+[build-url]: https://github.com/Chirayu-Tripathi/nl2query/actions/workflows/build.yaml
+[stars-image]: https://img.shields.io/github/stars/Chirayu-Tripathi/nl2query
+[stars-url]: https://github.com/Chirayu-Tripathi/nl2query

nl2query/mongoquery.py

@@ -166,7 +166,7 @@ def _load_model(self) -> object:
 
     def preprocess(self, db_schema: str, text: str) -> str:
         """Pre-Process the db_schema by removing new line and extra spaces, and creates a prompt for the model."""
-        db_schema = db_schema.replace("\n", "").replace("  ", "")
+        db_schema = db_schema.replace("\n", " ").replace("  ", "")
 
         prompt_template = f"""<s> 
         Task Description:

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nl2query"
-version = "0.1.7"
+version = "0.1.8"
 description = ""
 authors = ["Chirayu-Tripathi <chirayutripathi7@gmail.com>"]
 readme = "README.md"
@@ -12,6 +12,9 @@ regex = "^2023.6.3"
 transformers = "^4.39.0"
 pandas = "^1.5.3"
 peft = "^0.5.0"
+accelerate = "^0.29.3"
+bitsandbytes = "^0.43.1"
+einops = "^0.7.0"
 
 [tool.poetry.group.dev.dependencies]
 pre-commit = "^3.3.3"