Update app.py

app.py

@@ -1,4 +1,3 @@
-import os
 import gradio as gr
 import torch
 from langchain_huggingface import HuggingFaceEmbeddings
@@ -9,142 +8,78 @@ from langchain.chains import RetrievalQA
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from langchain_huggingface import HuggingFacePipeline
 
-# Configure GPU settings
-os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Using device: {device}")
-
-class CPSLChatbot:
-    def __init__(self):
-        self.initialize_components()
-
-    def initialize_components(self):
-        try:
-            # Load and process document
-            doc_loader = TextLoader("dataset.txt")
-            docs = doc_loader.load()
-            text_splitter = RecursiveCharacterTextSplitter(
-                chunk_size=1000, 
-                chunk_overlap=100
-            )
-            split_docs = text_splitter.split_documents(docs)
-
-            # Initialize embeddings and vector store
-            self.embeddings = HuggingFaceEmbeddings(
-                model_name="all-MiniLM-L6-v2",
-                model_kwargs={'device': device}
-            )
-            self.vectordb = FAISS.from_documents(split_docs, self.embeddings)
-
-            # Load model and tokenizer
-            model_name = "01-ai/Yi-Coder-9B-Chat"
-            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-            self.model = AutoModelForCausalLM.from_pretrained(
-                model_name,
-                device_map="auto",
-                torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-                trust_remote_code=True
-            )
-
-            # Set up QA pipeline
-            self.qa_pipeline = pipeline(
-                "text-generation",
-                model=self.model,
-                tokenizer=self.tokenizer,
-                max_new_tokens=750,
-                pad_token_id=self.tokenizer.eos_token_id,
-                device=0 if device == "cuda" else -1
-            )
-
-            # Initialize LangChain components
-            llm = HuggingFacePipeline(pipeline=self.qa_pipeline)
-            retriever = self.vectordb.as_retriever(search_kwargs={"k": 5})
-            self.qa_chain = RetrievalQA.from_chain_type(
-                retriever=retriever,
-                chain_type="stuff",
-                llm=llm,
-                return_source_documents=False
-            )
-            print("Initialization completed successfully")
-
-        except Exception as e:
-            print(f"Initialization error: {str(e)}")
-            raise
-
-    def preprocess_query(self, query):
-        if "script" in query.lower() or "code" in query.lower():
-            return f"Write a CPSL script: {query}"
-        return query
-
-    def clean_response(self, response):
-        result = response.get("result", "")
-        if "Answer:" in result:
-            return result.split("Answer:")[1].strip()
-        return result.strip()
-
-    def get_response(self, user_input):
-        try:
-            processed_query = self.preprocess_query(user_input)
-            raw_response = self.qa_chain.invoke({"query": processed_query})
-            return self.clean_response(raw_response)
-        except Exception as e:
-            return f"Error processing query: {str(e)}"
-
-def create_gradio_interface():
-    chatbot = CPSLChatbot()
-    
-    with gr.Blocks(title="CPSL Chatbot") as chat_interface:
-        gr.Markdown("# CPSL Chatbot with GPU Support")
-        gr.Markdown("Using Yi-Coder-9B-Chat model for CPSL script generation and queries")
-        
-        chat_history = gr.Chatbot(
-            value=[], 
-            elem_id="chatbot",
-            height=600
-        )
-        
-        with gr.Row():
-            user_input = gr.Textbox(
-                label="Your Message:",
-                placeholder="Type your message here...",
-                show_label=True,
-                elem_id="user-input"
-            )
-            send_button = gr.Button("Send", variant="primary")
-
-        def chat_response(user_message, history):
-            if not user_message:
-                return history, history
-            
-            bot_response = chatbot.get_response(user_message)
-            history.append((user_message, bot_response))
-            return history, history
-
-        send_button.click(
-            chat_response,
-            inputs=[user_input, chat_history],
-            outputs=[chat_history, chat_history],
-            api_name="chat"
-        )
-        
-        # Clear the input textbox after sending
-        send_button.click(lambda: "", None, user_input)
-        
-        # Also allow Enter key to send message
-        user_input.submit(
-            chat_response,
-            inputs=[user_input, chat_history],
-            outputs=[chat_history, chat_history],
-        )
-        user_input.submit(lambda: "", None, user_input)
-
-    return chat_interface
 
+# Load and process the document
+doc_loader = TextLoader("dataset.txt")
+docs = doc_loader.load()
+text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
+split_docs = text_splitter.split_documents(docs)
+
+# Create embeddings and vector store
+embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
+vectordb = FAISS.from_documents(split_docs, embeddings)
+
+# Load model and tokenizer
+model_name = "01-ai/Yi-Coder-9B-Chat"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name, 
+    device_map="auto", 
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32
+)
+
+# Set up the QA pipeline
+qa_pipeline = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    max_new_tokens=750,
+    pad_token_id=tokenizer.eos_token_id
+)
+
+llm = HuggingFacePipeline(pipeline=qa_pipeline)
+
+# Set up the retriever and QA chain
+retriever = vectordb.as_retriever(search_kwargs={"k": 5})
+qa_chain = RetrievalQA.from_chain_type(
+    retriever=retriever,
+    chain_type="stuff",
+    llm=llm,
+    return_source_documents=False
+)
+
+def preprocess_query(query):
+    if "script" in query or "code" in query.lower():
+        return f"Write a CPSL script: {query}"
+    return query
+
+def clean_response(response):
+    result = response.get("result", "")
+    if "Answer:" in result:
+        return result.split("Answer:")[1].strip()
+    return result.strip()
+
+def chatbot_response(user_input):
+    processed_query = preprocess_query(user_input)
+    raw_response = qa_chain.invoke({"query": processed_query})
+    return clean_response(raw_response)
+
+# Gradio interface
+with gr.Blocks() as chat_interface:
+    gr.Markdown("# CPSL Chatbot")
+    chat_history = gr.Chatbot(type='messages')
+    user_input = gr.Textbox(label="Your Message:")
+    send_button = gr.Button("Send")
+
+    def interact(user_message, history):
+        bot_reply = chatbot_response(user_message)
+        history.append({"role": "user", "content": user_message})
+        history.append({"role": "assistant", "content": bot_reply})
+        return history, history
+
+    send_button.click(interact, inputs=[user_input, chat_history], outputs=[chat_history, chat_history])
+
+# Launch the interface
 if __name__ == "__main__":
-    interface = create_gradio_interface()
-    interface.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=True,
-        enable_queue=True
-    )
+    chat_interface.launch()