Update app.py

app.py

@@ -205,36 +205,41 @@ class GradioInference:
         #     results["text"], max_length=150, min_length=30, do_sample=False
         # )
 
-        ########################## PRUEBA CHATGPT #################################
-        from langchain.chat_models import ChatOllama
-        from langchain.chains.llm import LLMChain
-        from langchain.prompts import PromptTemplate
-        from langchain.chains.combine_documents.stuff import StuffDocumentsChain
-        from langchain.document_loaders import WebBaseLoader
+        ########################## PRUEBA LLAMA2 #################################
+        from langchain import HuggingFacePipeline, PromptTemplate, LLMChain
         
-        # OPENAI_API_KEY = os.environ['OPENAI_API_KEY']
-        loader = WebBaseLoader("https://lilianweng.github.io/posts/2023-06-23-agent/")
-        
-        text = results["text"]
-        # Define prompt
-        prompt_template = """Write a concise summary of the following:
-        "{text}"
-        CONCISE SUMMARY:"""
-        prompt = PromptTemplate.from_template(prompt_template)
-        
-        # Define LLM chain
-        llm = ChatOllama(model_name="llama2:7b-chat")
-        llm_chain = LLMChain(llm=llm, prompt=prompt)
+        llm_model = "meta-llama/Llama-2-7b-chat-hf"
+
+        llm_tokenizer = AutoTokenizer.from_pretrained(model)
         
-        # Define StuffDocumentsChain
-        stuff_chain = StuffDocumentsChain(
-            llm_chain=llm_chain, document_variable_name="text"
+        pipeline = transformers.pipeline(
+            "text-generation", #task
+            model=llm_model,
+            tokenizer=llm_tokenizer,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            device_map="auto",
+            max_length=1000,
+            do_sample=True,
+            top_k=10,
+            num_return_sequences=1,
+            eos_token_id=tokenizer.eos_token_id
         )
-        
-        docs = loader.load()
-        sum = stuff_chain.run(docs)
 
-        ########################## FIN PRUEBA CHATGPT #################################
+        template = """
+              Write a concise summary of the following text delimited by triple backquotes.
+              ```{text}```
+              CONCISE SUMMARY:
+           """
+
+        prompt = PromptTemplate(template=template, input_variables=["text"])
+
+        llm_chain = LLMChain(prompt=prompt, llm=llm)
+
+        text = results["text"]
+
+        sum = llm_chain.run(text)
+        ########################## FIN PRUEBA LLAMA2 #################################
         
         #### Prueba
         # WHITESPACE_HANDLER = lambda k: re.sub('\s+', ' ', re.sub('\n+', ' ', k.strip()))