update model_operations.py for new llms

src/backend/model_operations.py

@@ -164,7 +164,7 @@ class SummaryGenerator:
         using_replicate_api = False
         replicate_api_models = ['snowflake', 'llama-3.1-405b']
         using_pipeline = False
-        pipeline_models = ['llama-3.1', 'phi-3-mini','falcon-7b', 'phi-3.5', 'mistral-nemo']
+        pipeline_models = ['llama-3.1', 'phi-3-mini','falcon-7b', 'phi-3.5', 'mistral-nemo', 'llama-3.3']
 
         for replicate_api_model in replicate_api_models:
             if replicate_api_model in self.model_id.lower():
@@ -222,6 +222,7 @@ class SummaryGenerator:
             print(result)
             return result
 
+        # Using Grok API
         elif 'grok' in self.model_id.lower(): # xai
             XAI_API_KEY = os.getenv("XAI_API_KEY")
             client = OpenAI(
@@ -241,6 +242,7 @@ class SummaryGenerator:
             print(result)
             return result
 
+        # Using Vertex AI API for Gemini models
         elif 'gemini' in self.model_id.lower():
             vertexai.init(project=os.getenv("GOOGLE_PROJECT_ID"), location="us-central1")
             model = GenerativeModel(
@@ -249,7 +251,7 @@ class SummaryGenerator:
             )
             generation_config = {
                 "temperature": 0,
-                "max_output_tokens": 250
+                "max_output_tokens": 500
             }
             safety_settings = [
                 SafetySetting(
@@ -277,6 +279,8 @@ class SummaryGenerator:
             result = response.text
             print(result)
             return result
+        
+        # Using Replicate API
         elif using_replicate_api:
             print("using replicate")
             if 'snowflake' in self.model_id.lower():
@@ -306,6 +310,7 @@ class SummaryGenerator:
             print(response)
             return response
 
+        # Using Anthropic API for Claude models
         elif 'claude' in self.model_id.lower(): # using anthropic api
             print('using Anthropic API')
             client = anthropic.Anthropic()
@@ -331,6 +336,7 @@ class SummaryGenerator:
             print(result)
             return result
 
+        # Using Cohere API
         elif 'command-r' in self.model_id.lower() or 'aya-expanse' in self.model_id.lower():
             co = cohere.ClientV2(os.getenv('COHERE_API_TOKEN'))
             response = co.chat(
@@ -345,6 +351,7 @@ class SummaryGenerator:
             print(result)
             return result
 
+        # Using MistralAI API
         elif 'mistral-large' in self.model_id.lower():
             api_key = os.environ["MISTRAL_API_KEY"]
             client = Mistral(api_key=api_key)
@@ -369,6 +376,7 @@ class SummaryGenerator:
             print(result)
             return result
 
+        # Using Deepseek API
         elif 'deepseek' in self.model_id.lower():
             client = OpenAI(api_key=os.getenv("DeepSeek_API_KEY"), base_url="https://api.deepseek.com")
             response = client.chat.completions.create(
@@ -385,20 +393,21 @@ class SummaryGenerator:
             print(result)
             return result
         
-        # Using HF API or download checkpoints
+        # Using HF pipeline or local checkpoints
         elif self.local_model is None and self.local_pipeline is None:
             if using_pipeline:
                 self.local_pipeline = pipeline(
                     "text-generation",
                     model=self.model_id,
                     tokenizer=AutoTokenizer.from_pretrained(self.model_id),
-                    torch_dtype=torch.bfloat16 if 'llama-3.2' in self.model_id.lower() else "auto",
+                    torch_dtype=torch.bfloat16 if 'llama-3.2' in self.model_id.lower() or 'llama-3.3' in self.model_id.lower() else "auto",
                     device_map="auto",
                     trust_remote_code=True
                 )
             else:
                 if 'ragamuffin' in self.model_id.lower():
                     self.tokenizer = AutoTokenizer.from_pretrained(os.path.join('/home/miaoran', self.model_id))
+
                 else:
                     self.tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf" if 'openelm' in self.model_id.lower() else self.model_id, trust_remote_code=True)
                 print("Tokenizer loaded")
@@ -420,7 +429,12 @@ class SummaryGenerator:
                 #     self.local_model = AutoModelForCausalLM.from_pretrained(os.path.join('/home/miaoran', self.model_id),
                 #                                  torch_dtype=torch.bfloat16, # forcing bfloat16 for now
                 #                                  attn_implementation="flash_attention_2")
-    
+                elif 'olmo' in self.model_id.lower():
+                    self.local_model = AutoModelForCausalLM.from_pretrained(self.model_id)#torch_dtype="auto"
+                
+                elif 'qwq-' in self.model_id.lower():
+                    self.local_model = AutoModelForCausalLM.from_pretrained(self.model_id, torch_dtype="auto", device_map="auto")
+
                 else:
                     self.local_model = AutoModelForCausalLM.from_pretrained(self.model_id, trust_remote_code=True, device_map="auto")#torch_dtype="auto"
                 # print(self.local_model.device)
@@ -435,7 +449,7 @@ class SummaryGenerator:
             ]
             outputs = self.local_pipeline(
                 messages,
-                max_new_tokens=250,
+                max_new_tokens=256,
                 # return_full_text=False,
                 do_sample=False
             )
@@ -445,6 +459,8 @@ class SummaryGenerator:
 
         elif self.local_model: # cannot call API. using local model / pipeline
             print('Using local model')
+
+            # Set appropriate prompt based on model document
             if 'gemma' in self.model_id.lower() or 'mistral-7b' in self.model_id.lower():
                 messages=[
                     # gemma-1.1, mistral-7b does not accept system role
@@ -478,29 +494,41 @@ class SummaryGenerator:
                     {"role": "system", "content": system_prompt}, 
                     {"role": "user", "content": user_prompt}
                 ]
-                prompt = self.tokenizer.apply_chat_template(messages,add_generation_prompt=True, tokenize=False)
-            # print(prompt)
-            # print('-'*50)
-            input_ids = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+                prompt = self.tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+           
+            # Tokenize inputs
+            if 'olmo' in self.model_id.lower():
+                input_ids = self.tokenizer([prompt], return_tensors='pt', return_token_type_ids=False)#.to(self.device)
+            elif 'qwq' in self.model_id.lower():
+                input_ids = self.tokenizer([prompt], return_tensors="pt").to(self.device)
+            else:
+                input_ids = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+            
+            # Generate outputs
             if 'granite' in self.model_id.lower():
                 self.local_model.eval()
                 outputs = self.local_model.generate(**input_ids, max_new_tokens=250)
+            elif 'olmo' in self.model_id.lower():
+                outputs = self.local_model.generate(**input_ids, max_new_tokens=250, do_sample=True, temperature=0.01)#top_k=50, top_p=0.95)
+            elif 'qwq' in self.model_id.lower():
+                outputs = self.local_model.generate(**input_ids, max_new_tokens=512, do_sample=True, temperature=0.01)
             else:
                 with torch.no_grad():
                     outputs = self.local_model.generate(**input_ids, do_sample=True, max_new_tokens=250, temperature=0.01)#, pad_token_id=self.tokenizer.eos_token_id
             if 'glm' in self.model_id.lower() or 'ragamuffin' in self.model_id.lower() or 'granite' in self.model_id.lower():
                 outputs = outputs[:, input_ids['input_ids'].shape[1]:]
-            elif 'qwen2-vl' in self.model_id.lower() or 'qwen2.5' in self.model_id.lower():
+            elif 'qwen2-vl' in self.model_id.lower() or 'qwen2.5' in self.model_id.lower() or 'qwq-' in self.model_id.lower():
                 outputs = [
                     out_ids[len(in_ids) :] for in_ids, out_ids in zip(input_ids.input_ids, outputs)
                 ]
-
+            
+            # Decode outputs
             if 'qwen2-vl' in self.model_id.lower():
                 result = self.processor.batch_decode(
                     outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False
                 )[0]
-            # elif 'granite' in self.model_id.lower():
-            #     result = self.tokenizer.batch_decode(outputs)[0]
+            elif 'olmo' in self.model_id.lower() or 'qwq' in self.model_id.lower():
+                result = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
             else:
                 result = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
 
@@ -512,9 +540,9 @@ class SummaryGenerator:
                 result = result.split(messages[-1]['content'])[1].strip()
             elif 'qwen2-vl' in self.model_id.lower() or 'qwen2.5' in self.model_id.lower():
                 pass
+            elif 'olmo' in self.model_id.lower():
+                result = result.split("<|assistant|>\n")[-1]
             else:
-                # print(prompt)
-                # print('-'*50)
                 result = result.replace(prompt.strip(), '')
             
             print(result)