utils optimised

utils.py

@@ -11,6 +11,7 @@ from nltk.corpus import stopwords
 from collections import deque
 from typing import Tuple
 import torch
+import streamlit as st
 
 # LangChain components
 from langchain_community.document_loaders import PyPDFLoader
@@ -25,10 +26,10 @@ from sentence_transformers import CrossEncoder
 from sklearn.metrics.pairwise import cosine_similarity
 
 # Initialize NLTK stopwords
-# nltk.download('stopwords')
-# stop_words = set(stopwords.words('english'))
-nltk.data.path.append('./nltk_data')  # Point to local NLTK data
-stop_words = set(nltk.corpus.stopwords.words('english'))
+nltk.download('stopwords')
+stop_words = set(stopwords.words('english'))
+# nltk.data.path.append('./nltk_data')  # Point to local NLTK data
+# stop_words = set(nltk.corpus.stopwords.words('english'))
 
 # mount
 import sys
@@ -38,7 +39,7 @@ sys.path.append('/mount/src/gen_ai_dev')
 DATA_PATH = "./Infy financial report/"
 DATA_FILES = ["INFY_2022_2023.pdf", "INFY_2023_2024.pdf"]
 EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
-LLM_MODEL = "gpt2" # "HuggingFaceH4/zephyr-7b-beta"  # "microsoft/phi-2"
+LLM_MODEL = "gpt2"  # Or "distilgpt2"  # Or "HuggingFaceH4/zephyr-7b-beta" or "microsoft/phi-2"
 
 # Environment settings
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
@@ -214,24 +215,26 @@ guard = SafetyGuard()
 # LLM Initialization
 # ------------------------------
 try:
-    tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL)
-    model = AutoModelForCausalLM.from_pretrained(
-        LLM_MODEL,
-        device_map="cpu",
-        torch_dtype=torch.float32
-    )
-
-    generator = pipeline(
-        "text-generation",
-        model=model,
-        tokenizer=tokenizer,
-        max_new_tokens=400,
-        do_sample=True,
-        temperature=0.3,
-        top_k=30,
-        top_p=0.9,
-        repetition_penalty=1.2
-    )
+    @st.cache_resource
+    def load_generator():
+        tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL)
+        model = AutoModelForCausalLM.from_pretrained(
+            LLM_MODEL,
+            device_map="auto",
+            torch_dtype=torch.float16,
+            load_in_4bit=True
+        )
+        return pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            max_new_tokens=100,
+            do_sample=False,
+            temperature=0.7,
+            top_k=0,
+            top_p=1,
+        )
+    generator = load_generator()
 except Exception as e:
     print(f"Error loading model: {e}")
     raise
@@ -268,24 +271,24 @@ Answer:"""
         print(f"\n\n[For Debug Only] Prompt: {prompt}\n\n")
         response = generator(prompt)[0]['generated_text']
         print(f"\n\n[For Debug Only] response: {response}\n\n")
-        
+
         clean_response = extract_final_response(response)
         clean_response = guard.filter_output(clean_response)
         print(f"\n\n[For Debug Only] clean_response: {clean_response}\n\n")
-        
+
         query_embed = embeddings.embed_query(query)
         print(f"\n\n[For Debug Only] query_embed: {query_embed}\n\n")
-        
+
         response_embed = embeddings.embed_query(clean_response)
         print(f"\n\n[For Debug Only] response_embed: {response_embed}\n\n")
-        
+
         confidence = cosine_similarity([query_embed], [response_embed])[0][0]
         print(f"\n\n[For Debug Only] confidence: {confidence}\n\n")
-        
-        memory.add_interaction(query, clean_response)       
-        
+
+        memory.add_interaction(query, clean_response)
+
         print(f"\n\n[For Debug Only] I'm Done \n\n")
         return clean_response, round(confidence, 2)
 
     except Exception as e:
-        return f"Error processing request: {e}", 0.0
+        return f"Error processing request: {e}", 0.0