Create app.py

app.py

@@ -0,0 +1,115 @@
+import os
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from datasets import concatenate_datasets, load_dataset
+import gc
+import gradio as gr
+from peft import PeftModel, PeftConfig
+from langchain.chains import RetrievalQA
+from langchain_community.vectorstores import FAISS
+from langchain.embeddings import OpenAIEmbeddings
+from langchain.docstore.document import Document
+from langchain.llms import HuggingFacePipeline
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.chains.question_answering import load_qa_chain
+from langchain.prompts import PromptTemplate
+import torch
+import random
+from langchain.document_loaders import WebBaseLoader
+from langchain.text_splitter import CharacterTextSplitter
+from langchain.memory import ConversationBufferMemory
+import requests
+import re
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Load Samsum dataset for generating questions
+train_dataset = load_dataset("samsum", split='train', trust_remote_code=True)
+val_dataset = load_dataset("samsum", split='validation', trust_remote_code=True)
+samsum_dataset = concatenate_datasets([train_dataset, val_dataset])
+
+model_name = "google/flan-t5-base"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+base_model = AutoModelForSeq2SeqLM.from_pretrained(model_name, torch_dtype=torch.bfloat16).to(device)
+rlhf_model_path = "raghav-gaggar/PEFT_RLHF_TextSummarizer"
+config = PeftConfig.from_pretrained(rlhf_model_path)
+ppo_model = PeftModel.from_pretrained(base_model, rlhf_model_path).to(device)
+merged_model = ppo_model.merge_and_unload().to(device)
+
+base_model.eval()
+ppo_model.eval()
+merged_model.eval()
+
+dialogsum_dataset = load_dataset("knkarthick/dialogsum", trust_remote_code=True)
+
+def format_dialogsum_as_document(example):
+    return Document(page_content=f"Dialogue:\n {example['dialogue']}\n\nSummary: {example['summary']}")
+
+# Create documents from DialogSum dataset
+documents = []
+for split in ['train', 'validation', 'test']:
+    documents.extend([format_dialogsum_as_document(example) for example in dialogsum_dataset[split]])
+
+# Split the documents into chunks
+text_splitter = CharacterTextSplitter(chunk_size=5200, chunk_overlap=0)
+docs = text_splitter.split_documents(documents)
+
+# Create embeddings and vector store for DialogSum documents
+embeddings = HuggingFaceEmbeddings(
+    model_name="sentence-transformers/all-MiniLM-L6-v2",
+    model_kwargs={"device": "cuda" if torch.cuda.is_available() else "cpu"},
+    encode_kwargs={"batch_size": 32}
+)
+
+vector_store = FAISS.from_documents(docs, embeddings)
+
+# Initialize retriever for DialogSum documents
+retriever = vector_store.as_retriever(search_kwargs={"k": 1})
+
+prompt_template = """
+Concisely summarize the dialogue in the end, like the example provided -
+
+Example -
+{context}
+
+Dialogue to be summarized:
+{question}
+
+Summary:"""
+
+PROMPT = PromptTemplate(
+    template=prompt_template, input_variables=["context", "question"]
+)
+
+# Create a Hugging Face pipeline
+summarization_pipeline = pipeline(
+    "summarization",
+    model=merged_model,
+    tokenizer=tokenizer,
+    max_length=150,
+    min_length=20,
+    do_sample=False,
+)
+
+# Wrap the pipeline in a LangChain LLM
+llm = HuggingFacePipeline(pipeline=summarization_pipeline)
+
+qa_chain = RetrievalQA.from_chain_type(
+    llm, retriever=retriever, chain_type_kwargs={"prompt": PROMPT}
+)
+
+# Function for Gradio interface
+def summarize_conversation(question):
+    result = qa_chain({"query": question})
+    return result["result"]
+    
+# Create Gradio interface
+iface = gr.Interface(
+    fn=summarize_conversation,
+    inputs=gr.Textbox(lines=10, label="Enter conversation here"),
+    outputs=gr.Textbox(label="Summary"),
+    title="Conversation Summarizer",
+    description="Enter a conversation, and the AI will provide a concise summary."
+)
+
+# Launch the app
+iface.launch()