Update llm.py

llm.py

@@ -1,4 +1,5 @@
 from langchain_huggingface.llms import HuggingFacePipeline
+from transformers import pipeline
 from langchain import HuggingFaceHub
 from huggingface_hub import login
 import random
@@ -19,12 +20,19 @@ text = """
     """
 
 print("loading model")
-model = HuggingFacePipeline.from_model_id(
-    model_id="mistralai/Mistral-7B-Instruct-v0.3",
-    task="text-generation",
-    pipeline_kwargs={"temperature":1, "max_length":1000,
-    "repetition_penalty":1.25, "max_new_tokens": 2000}
-)
+# model = HuggingFacePipeline.from_model_id(
+#     model_id="mistralai/Mistral-7B-Instruct-v0.3",
+#     task="text-generation",
+#     pipeline_kwargs={"temperature":1, "max_length":1000,
+#     "repetition_penalty":1.25, "max_new_tokens": 2000}
+# )
+
+model = pipeline("text-generation", 
+                     model="mistralai/Mistral-7B-v0.1", 
+                     tokenizer="mistralai/Mistral-7B-v0.1", 
+                     device_map="auto",  # Adjust if GPU is available
+                     torch_dtype="auto")
+
 print("loading model done!")
 
 print(model("if 2x = 4 + 10 what will be x"))