minor update for src/model_operations.py

src/backend/model_operations.py

@@ -162,7 +162,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']
+        pipeline_models = ['llama-3.1', 'phi-3-mini','falcon-7b', 'phi-3.5']
 
         for replicate_api_model in replicate_api_models:
             if replicate_api_model in self.model_id.lower():
@@ -375,12 +375,19 @@ class SummaryGenerator:
                     model=self.model_id,
                     model_kwargs={"torch_dtype": torch.bfloat16},
                     device_map="auto",
+                    trust_remote_code=True
                 )
             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")
-                self.local_model = AutoModelForCausalLM.from_pretrained(self.model_id, trust_remote_code=True, device_map="auto", torch_dtype="auto")
-                print(self.local_model.device)
+                if 'jamba' in self.model_id.lower():
+                    self.local_model = AutoModelForCausalLM.from_pretrained(self.model_id,
+                                             torch_dtype=torch.bfloat16,
+                                             attn_implementation="flash_attention_2",
+                                             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)
                 print("Local model loaded")
         
             
@@ -394,6 +401,8 @@ class SummaryGenerator:
             outputs = self.local_pipeline(
                 messages,
                 max_new_tokens=250,
+                temperature=0.0,
+                do_sample=False
             )
             result = outputs[0]["generated_text"][-1]['content']
             print(result)
@@ -435,8 +444,8 @@ class SummaryGenerator:
                 result = result.split("### Assistant:\n")[-1]
 
             else:
-                print(prompt)
-                print('-'*50)
+                # print(prompt)
+                # print('-'*50)
                 result = result.replace(prompt.strip(), '')
             
             print(result)