Update chatbot.py

chatbot.py

@@ -11,6 +11,10 @@ from langchain_community.llms import Ollama
 from langchain_core.messages import HumanMessage, AIMessage
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.prompts import MessagesPlaceholder
+import torch
+from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
 
 
 
@@ -43,12 +47,26 @@ def create_chain(chains, pdf_doc):
 
 
 def create_model():
-    
-    llm = HuggingFaceLLM.from_pretrained(
-        repo_id="google/flan-t5-base", 
-        temperature=1.0, 
-        max_new_tokens=250
-    )
+    tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
+
+    model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2",
+                                                device_map='auto',
+                                                torch_dtype=torch.float16,
+                                                use_auth_token=True,
+                                                load_in_8bit=True,
+                                                )
+    pipe = pipeline("text-generation",
+                    model=model,
+                    tokenizer= tokenizer,
+                    torch_dtype=torch.bfloat16,
+                    device_map="auto",
+                    max_new_tokens = 1024,
+                    do_sample=True,
+                    top_k=10,
+                    num_return_sequences=1,
+                    eos_token_id=tokenizer.eos_token_id
+                    )
+    llm = HuggingFacePipeline(pipeline=pipe, model_kwargs={'temperature':0})
     return llm