use llama2 for testing the model

.github/workflows/main.yml

@@ -5,40 +5,48 @@ name: Python application
 
 on:
   push:
-    branches: [ "main" ]
+    branches: ["main"]
   pull_request:
-    branches: [ "main" ]
+    branches: ["main"]
 
 permissions:
   contents: read
 
 jobs:
   build:
-
     runs-on: ubuntu-latest
 
     steps:
-    - name: Install Ollama
-      run: |
-        apt-get update && apt-get install -y git build-essential curl
-        curl https://ollama.ai/install.sh | sh
-        ollama serve&
-        ollama pull mistral
-    - uses: actions/checkout@v4
-    - name: Set up Python 3.10
-      uses: actions/setup-python@v3
-      with:
-        python-version: "3.10"
-    - name: Install Poetry
-      run: |
-        curl -sSL https://install.python-poetry.org | python3 -
-        echo "$HOME/.local/bin" >> $GITHUB_PATH
-    - name: Install dependencies with Poetry
-      run: |
-        poetry install
-    - name: Lint with ruff
-      run: |
-        poetry run ruff check
-    - name: Test with pytest
-      run: |
-        poetry run pytest --reruns 1 --reruns-delay 1
+      - name: Install Ollama
+        run: |
+          apt-get update && apt-get install -y git build-essential curl
+          curl https://ollama.ai/install.sh | sh
+          ollama serve&
+          ollama pull mistral
+          ollama pull llama2
+      - uses: actions/checkout@v4
+      - name: Set up Python 3.10
+        uses: actions/setup-python@v3
+        with:
+          python-version: "3.10"
+      - name: Install Poetry
+        run: |
+          curl -sSL https://install.python-poetry.org | python3 -
+          echo "$HOME/.local/bin" >> $GITHUB_PATH
+      - name: Install dependencies with Poetry
+        run: |
+          poetry install
+  lint:
+    needs: build
+    runs-on: ubuntu-latest
+    steps:
+      - name: Lint with ruff
+        run: |
+          poetry run ruff check
+  test:
+    needs: lint
+    runs-on: ubuntu-latest
+    steps:
+      - name: Test with pytest
+        run: |
+          poetry run pytest --reruns 1 --reruns-delay 1

src/generate_conversation.py

@@ -12,12 +12,14 @@ from langchain_experimental.tabular_synthetic_data.prompts import (
 
 examples = [
     {
-        "example": """"Hola, como estas?", "Bien, y tu?, "Muy bien, gracias, y tu?", "Yo tambien muy bien, gracias", "Que bueno, nos vemos luego!", "Adios!"]"""
+        "example": """"Hello, how are you?", "Great! you?", "Good, good thanks", "Great! See you later!", "See you!"]"""
     }
 ]
 
 
-def generate_conversation(topic: str, sentiment: str = "positive") -> str:
+def generate_conversation(
+    topic: str, sentiment: str = "positive", llm: str = "mistral"
+) -> str:
     """Generate a conversation about a topic with a given sentiment.
 
     Args:
@@ -37,7 +39,7 @@ def generate_conversation(topic: str, sentiment: str = "positive") -> str:
         ),
     )
 
-    model = Ollama(model="mistral")
+    model = Ollama(model=llm)
 
     generator = SyntheticDataGenerator(template=template, llm=model)
     results = generator.generate(

src/summarization.py

@@ -4,7 +4,7 @@ from langchain_community.document_loaders import TextLoader
 from langchain_community.llms import Ollama
 
 
-def topics_for_text(file_conv: str) -> str:
+def topics_for_text(file_conv: str, llm: str = "mistral") -> str:
     """Get the topics for a text.
 
     Args:
@@ -18,7 +18,7 @@ def topics_for_text(file_conv: str) -> str:
     TOPICS:"""
 
     prompt = PromptTemplate.from_template(prompt_template)
-    llm = Ollama(model="mistral")
+    llm = Ollama(model=llm)
 
     loader = TextLoader(file_conv)
     docs = loader.load()

tests/test_sentiment.py

@@ -44,7 +44,7 @@ def test_analyze_sentiment(sentiment):
         template="Generate a basic sentence with this sentiment: {sentiment}.",
         input_variables=["sentiment"],
     )
-    model = Ollama(model="mistral")
+    model = Ollama(model="llama2")
     output_parser = StrOutputParser()
 
     chain = prompt | model | output_parser

tests/{test_generate_conversation.py → test_sentiments_and_topics.py}

@@ -1,18 +1,21 @@
 import os
 import tempfile
 
+import pytest
+
 from src.generate_conversation import generate_conversation
 from src.sentiment import analyze_sentiment
 from src.summarization import topics_for_text
 
+SENTIMENTS = ["neutral", "sadness", "neutral", "joy"]
+TOPICS = ["ai", "soccer", "art", "science"]
 
-def test_generate_conversation_positive_sentiment():
-    # Define the input parameters
-    topic = "ai"
-    sentiment = "admiration"
 
+@pytest.mark.parametrize("sentiment", SENTIMENTS)
+@pytest.mark.parametrize("topic", TOPICS)
+def test_generate_conversation(sentiment, topic):
     # Call the function
-    conversation = generate_conversation(topic, sentiment)
+    conversation = generate_conversation(topic, sentiment, llm="llama2")
 
     new_sentiment = analyze_sentiment(conversation)
     # Assert that the conversation has a positive sentiment
@@ -28,5 +31,4 @@ def test_generate_conversation_positive_sentiment():
 
     new_topics = topics_for_text(temp_filepath)
     os.remove(temp_filepath)
-
-    assert topic in new_topics, "Topic is not in the index of new_topics"
+    assert any(topic in word for word in new_topics)

tests/test_summarization.py

@@ -1,23 +0,0 @@
-import os
-
-from src.summarization import topics_for_text
-
-
-def test_topics_for_text():
-    # Define the input parameters
-
-    file_conv = os.path.abspath(
-        os.path.join(
-            os.path.dirname(__file__),
-            "..",
-            "conversations",
-            "technology_1706560176_9889.txt",
-        )
-    )
-
-    # Call the function
-    topics = topics_for_text(file_conv)
-
-    # Assert the expected output
-    assert isinstance(topics, list)
-    assert "technology" in topics