update main.py

main.py

@@ -12,43 +12,26 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
-# 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 = "lamdao/lora-trained-xl-colab"
-# tokenizer = AutoTokenizer.from_pretrained(model_id)
-# model = AutoModelForCausalLM.from_pretrained(model_id)
-
-pipeline = pipeline(
-    "text-generation",
+pipe = pipeline(
+    "text2text-generation",
     model=model,
     tokenizer=tokenizer,
-    max_length=512
+    max_length=200,
+    temperature=0.8,
+    top_p=0.95,
+    repetition_penalty=1.15,
+    do_sample=True
 )
 
-# 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?'))
+local_llm = HuggingFacePipeline(pipeline=pipe)
 loader = PyPDFLoader('bipolar.pdf')
 # loader = TextLoader('info.txt')
 document = loader.load()
-text_spliter = CharacterTextSplitter(chunk_size=100, chunk_overlap=0)
+text_spliter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
 texts = text_spliter.split_documents(document)
 embedding = HuggingFaceInstructEmbeddings()
 docsearch = Chroma.from_documents(texts, embedding, persist_directory='db')
@@ -58,15 +41,13 @@ qa_chain = RetrievalQA.from_chain_type(llm=local_llm,
                                        chain_type="stuff",
                                        retriever=retriever,
                                        return_source_documents=True)
-# question = input('prompt: ')
-# result = qa_chain({'query': question})
-# print('result: ', result['result'])
+
 def gradinterface(query,history):
     result = qa_chain({'query': query})
     return result['result']
 
 
-demo = gr.ChatInterface(fn=gradinterface, title='OUR_BOT')
+demo = gr.ChatInterface(fn=gradinterface, title='OUR_OWN_BOT')
 
 if __name__ == "__main__":
     demo.launch(share=True)