Views
No views yet
1import os
2import torch
3from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
4
5MODEL_CONFIG = {
6 "director": {
7 "name": "Agnuxo/Qwen2-1.5B-Instruct_MOE_Director_16bit",
8 "task": "text-generation",
9 },
10 "programming": {
11 "name": "Qwen/Qwen2-1.5B-Instruct",
12 "task": "text-generation",
13 },
14 "biology": {
15 "name": "Agnuxo/Qwen2-1.5B-Instruct_MOE_BIOLOGY_assistant_16bit",
16 "task": "text-generation",
17 },
18 "mathematics": {
19 "name": "Qwen/Qwen2-Math-1.5B-Instruct",
20 "task": "text-generation",
21 }
22}
23
24
25KEYWORDS = {
26 "biology": ["cell", "DNA", "protein", "evolution", "genetics", "ecosystem", "organism", "metabolism", "photosynthesis", "microbiology", "célula", "ADN", "proteína", "evolución", "genética", "ecosistema", "organismo", "metabolismo", "fotosíntesis", "microbiología"],
27 "mathematics": ["Math" "mathematics", "equation", "integral", "derivative", "function", "geometry", "algebra", "statistics", "probability", "ecuación", "integral", "derivada", "función", "geometría", "álgebra", "estadística", "probabilidad"],
28 "programming": ["python", "java", "C++", "HTML", "scrip", "code", "Dataset", "API", "framework", "debugging", "algorithm", "compiler", "database", "CSS", "JSON", "XML", "encryption", "IDE", "repository", "Git", "version control", "front-end", "back-end", "API", "stack trace", "REST", "machine learning"]
29}
30
31class MOELLM:
32 def __init__(self):
33 self.current_expert = None
34 self.current_model = None
35 self.current_tokenizer = None
36 self.device = "cuda" if torch.cuda.is_available() else "cpu"
37 print(f"Using device: {self.device}")
38 self.load_director_model()
39
40 def load_director_model(self):
41 """Loads the director model."""
42 print("Loading director model...")
43 model_name = MODEL_CONFIG["director"]["name"]
44 self.director_tokenizer = AutoTokenizer.from_pretrained(model_name)
45 self.director_model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16).to(self.device)
46 self.director_pipeline = pipeline(
47 MODEL_CONFIG["director"]["task"],
48 model=self.director_model,
49 tokenizer=self.director_tokenizer,
50 device=self.device
51 )
52 print("Director model loaded.")
53
54 def load_expert_model(self, expert):
55 """Dynamically loads an expert model, releasing memory from the previous model."""
56 if expert not in MODEL_CONFIG:
57 raise ValueError(f"Unknown expert: {expert}")
58
59 if self.current_expert != expert:
60 print(f"Loading expert model: {expert}...")
61
62 # Free memory from the current model if it exists
63 if self.current_model:
64 del self.current_model
65 del self.current_tokenizer
66 torch.cuda.empty_cache()
67
68 model_config = MODEL_CONFIG[expert]
69 self.current_tokenizer = AutoTokenizer.from_pretrained(model_config["name"])
70 self.current_model = AutoModelForCausalLM.from_pretrained(model_config["name"], torch_dtype=torch.float16).to(self.device)
71 self.current_expert = expert
72
73 print(f"{expert.capitalize()} model loaded.")
74
75 return pipeline(
76 MODEL_CONFIG[expert]["task"],
77 model=self.current_model,
78 tokenizer=self.current_tokenizer,
79 device=self.device
80 )
81
82 def determine_expert_by_keywords(self, question):
83 """Determines the expert based on keywords in the question."""
84 question_lower = question.lower()
85 for expert, keywords in KEYWORDS.items():
86 if any(keyword in question_lower for keyword in keywords):
87 return expert
88 return None
89
90 def determine_expert(self, question):
91 """Determines which expert should answer the question."""
92 expert = self.determine_expert_by_keywords(question)
93 if expert:
94 print(f"Expert determined by keyword: {expert}")
95 return expert
96
97 prompt = f"Classify the following question into one of these categories: programming, biology, mathematics. Question: {question}\nCategory:"
98 response = self.director_pipeline(prompt, max_length=100, num_return_sequences=1)[0]['generated_text']
99 expert = response.split(":")[-1].strip().lower()
100 if expert not in MODEL_CONFIG:
101 expert = "director"
102 print(f"Redirecting question to: {expert}")
103 return expert
104
105 def generate_response(self, question, expert):
106 """Generates a response using the appropriate model."""
107 try:
108 model = self.load_expert_model(expert)
109 prompt = f"Answer the following question as an expert in {expert}: {question}\nAnswer:"
110 response = model(prompt, max_length=200, num_return_sequences=1)[0]['generated_text']
111 return response.split("Answer:")[-1].strip()
112 except Exception as e:
113 print(f"Error generating response: {str(e)}")
114 return "Sorry, there was an error processing your request. Please try again."
115
116 def chat_interface(self):
117 """Simple chat interface."""
118 print("Welcome to the MOE-LLM chat. Type 'exit' to quit.")
119 while True:
120 question = input("\nYou: ")
121 if question.lower() in ['exit', 'quit']:
122 break
123
124 try:
125 expert = self.determine_expert(question)
126 response = self.generate_response(question, expert)
127 print(f"\n{expert.capitalize()}: {response}")
128 except Exception as e:
129 print(f"Error in chat: {str(e)}")
130 print("Please try asking another question.")
131
132if __name__ == "__main__":
133 moe_llm = MOELLM()
134 moe_llm.chat_interface()
135
136