Update main.py

main.py

@@ -2,7 +2,7 @@ from langchain.chains import RetrievalQA, ConversationalRetrievalChain
 from langchain.vectorstores import Chroma
 from langchain.text_splitter import CharacterTextSplitter
 from langchain.document_loaders import DirectoryLoader, TextLoader,PyPDFLoader
-from transformers import pipeline
+from transformers import pipeline, AutoModelForCausalLM
 from langchain.llms import HuggingFacePipeline
 from langchain.embeddings import HuggingFaceInstructEmbeddings
 import gradio as gr
@@ -10,22 +10,41 @@ from InstructorEmbedding import INSTRUCTOR
 import torch
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
-model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
-pipe = pipeline(
-    "text2text-generation",
+# model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-mrpc")
+# tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-mrpc")
+
+
+#
+# tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
+#
+# model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
+
+model_id = "gpt2-medium"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id)
+
+pipeline = pipeline(
+    "text-generation",
     model=model,
     tokenizer=tokenizer,
-    max_length=512,
-    temperature=0.5,
-    top_p=0.95,
-    repetition_penalty=1.15
+    max_length=100
 )
 
-local_llm = HuggingFacePipeline(pipeline=pipe)
-print(local_llm('What is the capital of Syria?'))
+# local_llm = HuggingFacePipeline(pipeline=pipeline)
+# pipe = pipeline(
+#     "text2text-generation",
+#     model=model,
+#     tokenizer=tokenizer,
+#     max_length=512,
+#     temperature=0.5,
+#     top_p=0.95,
+#     repetition_penalty=1.15
+# )
 
+local_llm = HuggingFacePipeline(pipeline=pipeline)
+# print(local_llm('What is the capital of Syria?'))
 loader = PyPDFLoader('bipolar.pdf')
 # loader = TextLoader('info.txt')
 document = loader.load()