Polish Ollama support

Author: Skazza94

README.md

@@ -101,7 +101,7 @@ The experiment results will be stored in the directory named `results_low_level`
 ## Support for New Models
 We rely on LangChain to provide a common interface to access different model APIs.
 You can add new supported models in the `netconfeval/common/model_configs.py` file.
-We currently support OpenAI models (`'type': 'openai'`) and HuggingFace models (`'type': 'HF'`) through a custom LangChain-compatible class (`netconfeval/foundation/langchain/hf.py`).
+We currently support OpenAI models (`'type': 'openai'`), Ollama models (`'type': 'Ollama'`), and HuggingFace models (`'type': 'HF'`) through a custom LangChain-compatible class (`netconfeval/foundation/langchain/hf.py`).
 
 To add a model, just add a new Dict element to the `model_configurations` Dict, by providing a unique key for it.
 The new model key is then automatically visible using the `--model` command line parameter of the `.py` tests of the benchmarks.
@@ -131,6 +131,16 @@ model_configurations = {
 }
 ```
 
+### Ollama Models
+The Ollama model Dict contains the following keys:
+```python
+{
+    'type': 'Ollama', # The type of the model, in this case 'Ollama'
+    'model_name': 'llama3:8b-instruct-fp16', # The model name taken from Ollama library
+    'num_predict': 4096, # Max output length
+}
+```
+
 ### HuggingFace Models
 The HuggingFace model Dict contains the following keys:
 ```python
@@ -158,7 +168,7 @@ model_configurations = {
 ```
 
 ### Adding new model types
-Aside from adding OpenAI and HuggingFace models, it is also possible to add new model types (for example Gemini by Google).
+Aside from adding OpenAI, Ollama, and HuggingFace models, it is also possible to add new model types (for example Gemini by Google).
 
 We will continuously improve support for different APIs, but if you want to contribute:
 - Define a new `type`, coherent with the model types (e.g., `google` for Google models);
@@ -187,4 +197,4 @@ If you use NetConfEval, please cite our paper:
 ```
 
 ## Help
-If you have any questions regarding our code or the paper, you can contact [Changjie Wang](https://www.kth.se/profile/changjie) (changjie at kth.se) and/or [Mariano Scazzariello](https://www.kth.se/profile/marianos) (mariano.scazzariello at ri.se).
+If you have any questions regarding our code or the paper, you can contact [Changjie Wang](https://www.kth.se/profile/changjie) (changjie at kth.se) and/or [Mariano Scazzariello](https://www.ri.se/en/person/mariano-scazzariello) (mariano.scazzariello at ri.se).

netconfeval/common/model_configs.py

@@ -106,11 +106,10 @@ def _build_mistral_lite_prompt(messages):
 
     return prompt
 
+
 def _build_qwen2_prompt(messages):
-    
     start_turn = "<|im_start|>"
     end_turn = "<|im_end|>\n"
- 
 
     conversation = []
     for index, message in enumerate(messages):
@@ -122,14 +121,14 @@ def _build_qwen2_prompt(messages):
         content = content.strip()
 
         if role.lower() in ['user', 'system', 'assistant']:
-            conversation.append(start_turn + role.lower() + '\n' + content + end_turn) 
+            conversation.append(start_turn + role.lower() + '\n' + content + end_turn)
         else:
             raise ValueError(f"Unexpected role: {role}")
 
-   # Assemble the prompt with the start and end tokens and start the turn of assistant to prime the generation process
-    prompt = ' '.join(conversation) + start_turn + "assistant\n" 
+    # Assemble the prompt with the start and end tokens and start the turn of assistant to prime the generation process
+    prompt = ' '.join(conversation) + start_turn + "assistant\n"
 
-    return prompt 
+    return prompt
 
 
 def _build_llama3_prompt(messages):
@@ -148,18 +147,17 @@ def _build_llama3_prompt(messages):
         content = content.strip()
 
         if role.lower() in ['user', 'system', 'assistant']:
-            conversation.append(start_role + role.lower() + end_role + content + end_turn ) # Example: <|begin_of_text|><|start_header_id|>system<|end_header_id|>
+            # Example: <|begin_of_text|><|start_header_id|>system<|end_header_id|>
+            conversation.append(start_role + role.lower() + end_role + content + end_turn)
         else:
             raise ValueError(f"Unexpected role: {role}")
 
-   # Assemble the prompt with the start and end tokens and start the turn of assistant to prime the generation process
-    prompt = start_prompt + ' '.join(conversation)  + '\n' + start_role + "assistant" + end_role
-    print(prompt)
+    prompt = start_prompt + ' '.join(conversation) + '\n' + start_role + "assistant" + end_role
+
     return prompt
 
 
 def get_model_instance(model_name: str) -> Any:
-
     if model_configurations[model_name]['type'] == 'HF':
         from netconfeval.foundation.langchain.chat_models.hf import ChatHF
 
@@ -169,15 +167,16 @@ def get_model_instance(model_name: str) -> Any:
             use_quantization=model_configurations[model_name]['use_quantization'],
             prompt_func=model_configurations[model_name]['prompt_builder'],
         )
-    
+
     elif model_configurations[model_name]['type'] == 'Ollama':
         from langchain_community.llms import Ollama
 
-        return Ollama(model = model_configurations[model_name]['model_name'],
-                      num_predict = model_configurations[model_name]['num_predict'],
-                      num_gpu=-1
-                     )
-        
+        return Ollama(
+            model=model_configurations[model_name]['model_name'],
+            num_predict=model_configurations[model_name]['num_predict'],
+            num_gpu=-1
+        )
+
     elif model_configurations[model_name]['type'] == 'openai':
         from langchain_openai import ChatOpenAI
 
@@ -190,6 +189,7 @@ def get_model_instance(model_name: str) -> Any:
 
 
 model_configurations = {
+    # OpenAI Models
     'gpt-3.5-turbo': {
         'model_name': 'gpt-3.5-turbo',
         'type': 'openai',
@@ -226,74 +226,59 @@ def get_model_instance(model_name: str) -> Any:
             'seed': 5000,
         }
     },
-
-    #### Start Ollama models ####
-
+    # Ollama Models
     'llama3.1-ollama': {
         'type': 'Ollama',
         'model_name': 'llama3.1:8b-instruct-fp16',
         'num_predict': 4096
     },
-
-    'llama3-ollama': {
+    'llama3.1-4bit-ollama': {
         'type': 'Ollama',
-        'model_name': 'llama3:8b-instruct-fp16',
+        'model_name': 'llama3.1:latest',
         'num_predict': 4096
     },
-
-    'neural-chat-ollama': {
+    'llama3-ollama': {
         'type': 'Ollama',
-        'model_name': 'neural-chat:7b-v3.3-fp16',
+        'model_name': 'llama3:8b-instruct-fp16',
         'num_predict': 4096
     },
-
-    # 4-bit quantization version 
-
-    'llama3.1-4bit-ollama': {
+    'llama3-4bit-ollama': {
         'type': 'Ollama',
-        'model_name': 'llama3.1:latest',
+        'model_name': 'llama3:latest',
         'num_predict': 4096
     },
-
-    'llama3-4bit-ollama': {
+    'neural-chat-ollama': {
         'type': 'Ollama',
-        'model_name': 'llama3:latest',
+        'model_name': 'neural-chat:7b-v3.3-fp16',
         'num_predict': 4096
     },
-      
     'neural-chat-4bit-ollama': {
         'type': 'Ollama',
         'model_name': 'neural-chat:latest',
         'num_predict': 4096
     },
-
-    #### End Ollama models #####
-    
-
+    # HuggingFace Models
     'qwen2.5-7b-instruct': {
         'type': 'HF',
-        'model_name':'Qwen/Qwen2.5-7B-Instruct',
+        'model_name': 'Qwen/Qwen2.5-7B-Instruct',
         'prompt_builder': _build_qwen2_prompt,
         'max_length': 4096,
         'use_quantization': False
     },
-
     'llama3-8b-instruct': {
         'type': 'HF',
         'model_name': 'meta-llama/Meta-Llama-3-8B-Instruct',
         'prompt_builder': _build_llama3_prompt,
         'max_length': 4096,
         'use_quantization': False
     },
-
     'llama3.1-8b-instruct': {
         'type': 'HF',
         'model_name': 'meta-llama/Meta-Llama-3.1-8B-Instruct',
         'prompt_builder': _build_llama3_prompt,
         'max_length': 4096,
         'use_quantization': False
     },
-
     'llama2-7b-chat': {
         'model_name': 'meta-llama/Llama-2-7b-chat-hf',
         'prompt_builder': _build_llama2_prompt,

netconfeval/step_1_formal_spec_conflict_detection.py

@@ -143,7 +143,7 @@ def main(args: argparse.Namespace) -> None:
 
                     skip_compare = False
                     start_time = time.time()
-                    if model_configurations[args.model]['type'] in ['HF','Ollama']:
+                    if model_configurations[args.model]['type'] in ['HF', 'Ollama']:
                         # Combine all system prompts with a new line separator
                         combined_system_prompt = f"{SETUP_PROMPT}\n{FUNCTION_PROMPT}\n{ASK_FOR_RESULT_PROMPT}"
                         messages = [

netconfeval/step_1_formal_spec_conflict_distance.py

@@ -137,7 +137,7 @@ def main(args: argparse.Namespace) -> None:
 
                             skip_compare = False
                             start_time = time.time()
-                            if model_configurations[args.model]['type']  in ['HF','Ollama']:
+                            if model_configurations[args.model]['type'] in ['HF', 'Ollama']:
                                 # Combine all system prompts with a new line separator
                                 combined_system_prompt = f"{SETUP_PROMPT}\n{FUNCTION_PROMPT}\n{ASK_FOR_RESULT_PROMPT}"
                                 messages = [

netconfeval/step_1_formal_spec_translation.py

@@ -135,7 +135,7 @@ def main(args: argparse.Namespace) -> None:
 
                     skip_compare = False
                     start_time = time.time()
-                    if model_configurations[args.model]['type'] in ['HF','Ollama']:
+                    if model_configurations[args.model]['type'] in ['HF', 'Ollama']:
                         # Combine all system prompts with a new line separator
                         combined_system_prompt = f"{SETUP_PROMPT}\n{FUNCTION_PROMPT}\n{ASK_FOR_RESULT_PROMPT}"
                         messages = [

netconfeval/step_2_code_gen.py

@@ -126,7 +126,7 @@ def main(args: argparse.Namespace) -> None:
                     w = csv.DictWriter(f, result_row.keys())
                     w.writeheader()
 
-                if model_configurations[args.model]['type'] in ['HF','Ollama']:
+                if model_configurations[args.model]['type'] in ['HF', 'Ollama']:
                     combined_system_prompt = f"{SETUP_PROMPT}\n{ASK_FOR_CODE_PROMPT}"
                     combined_human_prompt = f"{INPUT_OUTPUT_PROMPT}\n{INSTRUCTION_PROMPT}\n{{input}}"
                     if with_feedback: