Second Commit

app.py

@@ -1,82 +1,146 @@
-import gradio as gr
+import os
+from typing import List
+from langchain.vectorstores.pgvector import PGVector
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.llms import HuggingFacePipeline
+from langchain.prompts import PromptTemplate
+from langchain.chains import RetrievalQA
+from sentence_transformers import SentenceTransformer
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import psycopg2
+from psycopg2.extras import execute_values
+import numpy as np
 
-# Load model and tokenizer
-model = AutoModelForCausalLM.from_pretrained("deepseek-ai/DeepSeek-R1", trust_remote_code=True)
-tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/DeepSeek-R1", trust_remote_code=True)
-
-# Move model to GPU if available
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model = model.to(device)
-
-def generate_response(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    # Format conversation history
-    messages = [{"role": "system", "content": system_message}]
+class RAGPipeline:
+    def __init__(self):
+        # Database connection string
+        self.connection_string = "postgresql://Data_owner:JsxygNDC15IO@ep-cool-hill-a5k13m05-pooler.us-east-2.aws.neon.tech/Data?sslmode=require"
+        
+        # Initialize embedding model
+        self.embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+        
+        # Initialize LLM
+        self.llm_model = AutoModelForCausalLM.from_pretrained(
+            "deepseek-ai/DeepSeek-R1",
+            trust_remote_code=True,
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+        )
+        self.llm_tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/DeepSeek-R1", trust_remote_code=True)
+        
+        # Move model to GPU if available
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.llm_model = self.llm_model.to(self.device)
+        
+        # Initialize prompt template
+        self.prompt_template = """
+        Use the following context to answer the question. If you cannot answer the question based on the context, say so.
+        
+        Context: {context}
+        
+        Question: {question}
+        
+        Answer: Let me help you with that.
+        """
     
-    for user_msg, assistant_msg in history:
-        if user_msg:
-            messages.append({"role": "user", "content": user_msg})
-        if assistant_msg:
-            messages.append({"role": "assistant", "content": assistant_msg})
+    def generate_embedding(self, text: str) -> List[float]:
+        """Generate embeddings for input text."""
+        embedding = self.embedding_model.encode(text)
+        return embedding.tolist()
     
-    messages.append({"role": "user", "content": message})
+    def similarity_search(self, query_embedding: List[float], top_k: int = 3) -> List[dict]:
+        """Perform similarity search in PostgreSQL using vector comparison."""
+        with psycopg2.connect(self.connection_string) as conn:
+            with conn.cursor() as cur:
+                # Convert query embedding to PostgreSQL array format
+                embedding_array = np.array(query_embedding)
+                
+                # Perform similarity search using cosine distance
+                query = """
+                SELECT text, title, url, 
+                       1 - (vector <=> %s) as similarity
+                FROM bents
+                ORDER BY vector <=> %s
+                LIMIT %s;
+                """
+                cur.execute(query, (embedding_array.tolist(), embedding_array.tolist(), top_k))
+                results = cur.fetchall()
+                
+                # Format results
+                similar_docs = [
+                    {
+                        'text': row[0],
+                        'title': row[1],
+                        'url': row[2],
+                        'similarity': row[3]
+                    }
+                    for row in results
+                ]
+                
+                return similar_docs
     
-    # Convert messages to model input format
-    conversation = ""
-    for msg in messages:
-        role = msg["role"]
-        content = msg["content"]
-        if role == "system":
-            conversation += f"System: {content}\n"
-        elif role == "user":
-            conversation += f"User: {content}\n"
-        elif role == "assistant":
-            conversation += f"Assistant: {content}\n"
+    def generate_response(self, query: str, context: str) -> str:
+        """Generate response using the LLM."""
+        # Format prompt with context and query
+        prompt = self.prompt_template.format(context=context, question=query)
+        
+        # Tokenize input
+        inputs = self.llm_tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048).to(self.device)
+        
+        # Generate response
+        with torch.no_grad():
+            outputs = self.llm_model.generate(
+                **inputs,
+                max_new_tokens=512,
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.95,
+                pad_token_id=self.llm_tokenizer.eos_token_id,
+            )
+        
+        # Decode and return response
+        response = self.llm_tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
+        return response.strip()
     
-    # Tokenize input
-    inputs = tokenizer(conversation, return_tensors="pt", truncation=True, max_length=2048).to(device)
+    def process_query(self, query: str) -> dict:
+        """Process user query through the complete RAG pipeline."""
+        # Generate embedding for query
+        query_embedding = self.generate_embedding(query)
+        
+        # Perform similarity search
+        similar_docs = self.similarity_search(query_embedding)
+        
+        # Combine relevant context
+        context = "\n".join([doc['text'] for doc in similar_docs])
+        
+        # Generate response
+        response = self.generate_response(query, context)
+        
+        return {
+            'query': query,
+            'response': response,
+            'similar_documents': similar_docs
+        }
+
+# Example usage
+def main():
+    # Initialize RAG pipeline
+    rag = RAGPipeline()
     
-    # Generate response
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=max_tokens,
-            do_sample=True,
-            temperature=temperature,
-            top_p=top_p,
-            pad_token_id=tokenizer.eos_token_id,
-        )
+    # Example query
+    query = "What are some woodworking tools that every woodworker should have?"
     
-    # Decode and return response
-    response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
-    return response.strip()
-
-# Create Gradio interface
-demo = gr.ChatInterface(
-    generate_response,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-    title="DeepSeek-R1 Chat",
-    description="A chat interface powered by the DeepSeek-R1 language model.",
-)
+    # Process query
+    result = rag.process_query(query)
+    
+    # Print results
+    print("\nQuery:", result['query'])
+    print("\nResponse:", result['response'])
+    print("\nRelevant Documents:")
+    for doc in result['similar_documents']:
+        print(f"\nTitle: {doc['title']}")
+        print(f"URL: {doc['url']}")
+        print(f"Similarity Score: {doc['similarity']:.4f}")
 
 if __name__ == "__main__":
-    demo.launch()
+    main()

requirements.txt

@@ -1 +1,23 @@
-huggingface_hub==0.25.2
+huggingface_hub==0.25.2
+# Core dependencies
+langchain>=0.1.0
+torch>=2.0.0
+transformers>=4.36.0
+sentence-transformers>=2.2.2
+
+# Database
+psycopg2-binary>=2.9.9
+pgvector>=0.2.3
+
+# Data processing
+numpy>=1.24.0
+pandas>=2.0.0
+
+# Deep learning
+accelerate>=0.24.0
+bitsandbytes>=0.41.0
+safetensors>=0.4.0
+
+# Utilities
+tqdm>=4.65.0
+python-dotenv>=1.0.0