Spaces:
Running
Running
| import os | |
| import re | |
| from ctransformers import AutoModelForCausalLM | |
| MODEL_REPO = "TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF" | |
| MODEL_FILE = "tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf" | |
| MODEL_TYPE = "llama" | |
| PROFANITY_LIST = ['damn', 'hell', 'crap', 'stupid', 'idiot', 'dumb', 'suck'] | |
| OFF_TOPIC_KEYWORDS = ['weather', 'sports', 'politics', 'religion', 'movie', 'game', 'recipe', 'joke', 'story', 'music', 'celebrity'] | |
| class LLMEngine: | |
| def __init__(self): | |
| self.model = None | |
| self.context = "" | |
| self.load_context() | |
| def load_context(self): | |
| try: | |
| context_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "data", "company_context.txt") | |
| if os.path.exists(context_path): | |
| with open(context_path, "r", encoding="utf-8") as f: | |
| self.context = f.read() | |
| except Exception as e: | |
| print(f"Error loading context: {e}") | |
| def load_model(self): | |
| if self.model is None: | |
| print("Loading LLM...") | |
| try: | |
| self.model = AutoModelForCausalLM.from_pretrained(MODEL_REPO, model_file=MODEL_FILE, model_type=MODEL_TYPE, context_length=2048, gpu_layers=0) | |
| print("LLM Loaded.") | |
| except Exception as e: | |
| print(f"Failed to load LLM: {e}") | |
| def check_profanity(self, text): | |
| for word in PROFANITY_LIST: | |
| if re.search(r'\b' + word + r'\b', text.lower()): | |
| return True | |
| return False | |
| def check_off_topic(self, text): | |
| for keyword in OFF_TOPIC_KEYWORDS: | |
| if keyword in text.lower(): | |
| return True | |
| return False | |
| def generate_response(self, user_query): | |
| if self.check_profanity(user_query): | |
| return "I'm here to assist with business operations. Please keep our conversation professional." | |
| if self.check_off_topic(user_query): | |
| return "I specialize in demand forecasting, inventory management, and order processing. How can I help you with these?" | |
| greetings = ["hi", "hello", "hey", "greetings", "good morning", "good afternoon", "good evening"] | |
| if user_query.lower().strip().strip("!.,?") in greetings: | |
| return "Hello! I'm BMS AI Assistant. I can help you with:\n• Demand Forecasting\n• Inventory Checks\n• Supplier Information\n• Order Requisitions\n• PDF Reports\n\nWhat would you like to know?" | |
| if any(w in user_query.lower() for w in ['capabilities', 'what can you do', 'help me', 'functions']): | |
| return "I can assist you with:\n1. Demand Forecasting\n2. Inventory Management\n3. Supplier Information\n4. Order Processing\n5. PDF Reports\n\nTry asking: 'Forecast for BMS0015'" | |
| if any(w in user_query.lower() for w in ['who made you', 'who developed', 'who created', 'who built']): | |
| return "I'm BMS AI Assistant, developed to help you manage inventory and forecast demand efficiently." | |
| if any(p in user_query.lower() for p in ['how are you', 'how do you do', 'how is it going']): | |
| return "I'm functioning well and ready to assist you! How can I help with your inventory or forecasting needs today?" | |
| if self.model is None: | |
| self.load_model() | |
| if self.model is None: | |
| return "I'm sorry, I couldn't load my language model. Please try asking about specific items." | |
| system_prompt = ( | |
| "You are BMS AI Assistant, a helpful and professional AI for Business Management Systems. " | |
| "Your goal is to assist users with demand forecasting, inventory management, and supplier information. " | |
| "Answer questions based ONLY on the provided context. If the answer is not in the context, politely state that you don't have that information. " | |
| "Be concise but friendly." | |
| ) | |
| full_prompt = f"<|system|>\n{system_prompt}\n\nContext:\n{self.context}</s>\n<|user|>\n{user_query}</s>\n<|assistant|>" | |
| try: | |
| response = self.model( | |
| full_prompt, | |
| max_new_tokens=150, | |
| temperature=0.1, | |
| repetition_penalty=1.1, | |
| top_k=40, | |
| top_p=0.9, | |
| stop=["</s>", "<|user|>", "<|system|>"] | |
| ) | |
| if "<|system|>" in response: | |
| response = response.split("<|system|>")[0] | |
| return response.strip() | |
| except Exception as e: | |
| return f"Error generating response: {e}" | |
| llm = LLMEngine() | |