refactor

.gitignore

@@ -1 +1,3 @@
 __pycache__
+.coverage
+*.log

app/interface.py

@@ -1,4 +1,5 @@
 import gradio as gr
+
 from .pipeline import Pipeline
 from .processing import process_file, process_text_input
 

app/models.py

@@ -0,0 +1,46 @@
+import json
+import csv
+from io import StringIO
+
+from pydantic import BaseModel, validator, ValidationError
+
+class Card(BaseModel):
+    question: str
+    answer: str
+
+class Message(BaseModel):
+    role: str
+    content: list[Card]
+
+    @validator('content', pre=True)
+    def parse_content(cls, v):
+        if isinstance(v, str):
+            try:
+                content_list = json.loads(v)
+                return content_list
+            except json.JSONDecodeError as e:
+                raise ValueError(f"Error decoding 'content' JSON: {e}") from e
+        return v
+
+    def content_to_json(self) -> str:
+        return json.dumps([card.dict() for card in self.content], indent=2)
+
+    def content_to_csv(self) -> str:
+        """
+        Converts the content of the Message instance into a CSV string.
+        """
+        output = StringIO()
+        # Step 2: Create a CSV writer - windows style is the default, so set Unix-style line endings
+        writer = csv.writer(output, lineterminator='\n')
+        writer.writerow(["Question", "Answer"])
+        for card in self.content:
+            writer.writerow([card.question, card.answer])
+        csv_content = output.getvalue()
+        output.close()
+        return csv_content
+
+class PydanticEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, BaseModel):
+            return obj.dict()
+        return super().default(obj)

app/pipeline.py

@@ -1,48 +1,14 @@
 from io import StringIO
-import csv
 import json
 import logging
 
 import torch
 from transformers import pipeline
-from pydantic import BaseModel, ValidationError, validator
 
-logger = logging.getLogger(__name__)
-
-class Card(BaseModel):
-    question: str
-    answer: str
-
-class Message(BaseModel):
-    role: str
-    content: list[Card]
-
-    @validator('content', pre=True)
-    def parse_content(cls, v):
-        if isinstance(v, str):
-            try:
-                content_list = json.loads(v)
-                return content_list
-            except json.JSONDecodeError as e:
-                raise ValueError(f"Error decoding 'content' JSON: {e}") from e
-        return v
-
-    def content_to_json(self) -> str:
-        return json.dumps([card.dict() for card in self.content], indent=2)
-
-    def content_to_csv(self) -> str:
-        output = StringIO()
-        writer = csv.writer(output)
-        writer.writerow(['Question', 'Answer'])  # CSV Header
-        for card in self.content:
-            writer.writerow([card.question, card.answer])
-        return output.getvalue()
+from .models import Card, Message, ValidationError
 
-class PydanticEncoder(json.JSONEncoder):
-    def default(self, obj):
-        if isinstance(obj, BaseModel):
-            return obj.dict()
-        return super().default(obj)
+logger = logging.getLogger(__name__)
+logging.basicConfig(filename="pipeline.log", level=logging.INFO, datefmt="%Y-%m-%d %H:%M:%S")
 
 class Pipeline:
     def __init__(self, model_name: str = "Qwen/Qwen2.5-7B-Instruct"):
@@ -53,6 +19,7 @@ class Pipeline:
             device_map="auto"
         )
         self.device = self._determine_device()
+        logger.info(f"device type: {self.device}")
         self.messages = [
             {"role": "system", "content": """You are an expert flashcard creator. You always include a single knowledge item per flashcard.
             - You ALWAYS include a single knowledge item per flashcard.
@@ -77,46 +44,9 @@ class Pipeline:
         )[0]["generated_text"][-1]
         return response_message
 
-    def format_flashcards(self, output_format: str, response: str) -> str:
-        output = ""
-        try :
-            message = parse_message(response)
-            logger.debug("after parse_obj_as")
-        except ValidationError as e:
-            raise e
-        if output_format.lower() == "json":
-            output = message.content_to_json()
-        elif output_format.lower() == "csv":
-            output = message.content_to_csv()
-        return output
-
     def generate_flashcards(self, output_format: str, content: str) -> str:
         response = self.extract_flashcards(content)
-        return self.format_flashcards(output_format, response)
-
-    def parse_message(self, input_dict: dict[str, any]) -> Message:
-        try:
-            # Extract the role
-            role = input_dict['role']
-
-            # Parse the content
-            content = input_dict['content']
-            # If content is a string, try to parse it as JSON
-            if isinstance(content, str):
-                content = content.strip()
-                content = json.loads(content)
-
-            # Create Card objects from the content
-            cards = [Card(**item) for item in content]
-
-            # Create and return the Message object
-            return Message(role=role, content=cards)
-        except json.JSONDecodeError as e:
-            raise ValueError(f"Invalid JSON in content: {str(e)}")
-        except ValidationError as e:
-            raise ValueError(f"Validation error: {str(e)}")
-        except KeyError as e:
-            raise ValueError(f"Missing required key: {str(e)}")
+        return format_flashcards(output_format, response)
 
     def _determine_device(self):
         if torch.cuda.is_available():
@@ -144,8 +74,11 @@ def parse_message(input_dict: dict[str, any]) -> Message:
         # Create and return the Message object
         return Message(role=role, content=cards)
     except json.JSONDecodeError as e:
+        logger.error(f"Invalid JSON in content: {str(e)}")
         raise ValueError(f"Invalid JSON in content: {str(e)}")
     except ValidationError as e:
+        logger.error(f"Validation error: {str(e)}")
         raise ValueError(f"Validation error: {str(e)}")
     except KeyError as e:
+        logger.error(f"Missing required key: {str(e)}")
         raise ValueError(f"Missing required key: {str(e)}")

app/processing.py

@@ -22,47 +22,6 @@ def read_text_file(file_path: str) -> str:
     except Exception as e:
         raise ValueError(f"Error reading text file: {str(e)}")
 
-def format_prompt(output_format: str) -> str:
-    """
-    Formats the prompt based on the output type.
-    """
-    if output_format.lower() == "json":
-        return """You only respond in JSON format. Follow the example below.
-
-    EXAMPLE:
-    [
-        {"question": "What is AI?", "answer": "Artificial Intelligence."},
-        {"question": "What is ML?", "answer": "Machine Learning."}
-    ]
-    """
-    elif output_format.lower() == "csv":
-        return """You only respond with cards in CSV format. Follow the example below.
-
-    EXAMPLE:
-        "What is AI?", "Artificial Intelligence."
-        "What is ML?", "Machine Learning."
-    """
-
-# def extract_flashcards(text: str, output_format: str, pipeline: str) -> str:
-#     """
-#     Extracts flashcards from the input text using the LLM and formats them in CSV or JSON.
-#     """
-#     prompt = f"""You are an expert flashcard creator. You always include a single knowledge item per flashcard.
-
-#     {format_prompt(output_format)}
-
-
-#     Extract flashcards from the user's text:
-
-#     {text}
-
-#     Do not include the prompt or any other unnecessary information in the flashcards.
-#     Do not include triple ticks (```) or any other code blocks in the flashcards.
-#     """
-#     # TODO:
-#     response = pipeline.generate_flashcards("json", prompt)
-#     return response
-
 def process_file(file_obj, output_format: str, pipeline) -> str:
     """
     Processes the uploaded file based on its type and extracts flashcards.
@@ -89,3 +48,17 @@ def process_text_input(output_format: str, input_text: str) -> str:
 
     flashcards = pipeline.generate_flashcards(output_format, input_text)
     return flashcards
+
+
+def format_flashcards(self, output_format: str, response: str) -> str:
+    output = ""
+    try :
+        message = parse_message(response)
+        logger.debug("after parse_obj_as")
+    except ValidationError as e:
+        raise e
+    if output_format.lower() == "json":
+        output = message.content_to_json()
+    elif output_format.lower() == "csv":
+        output = message.content_to_csv()
+    return output

tests/conftest.py

@@ -1,6 +1,6 @@
 import pytest
 from unittest.mock import Mock
-from app.pipeline import LanguageModel
+from app.pipeline import Pipeline
 
 @pytest.fixture
 def pipeline():
@@ -8,7 +8,7 @@ def pipeline():
     Fixture to provide a mocked LanguageModel instance.
     """
     # Create a mock instance of LanguageModel
-    lm = Mock(spec=LanguageModel)
+    lm = Mock(spec=Pipeline)
     # Mock the generate_flashcards method
     lm.generate_flashcards.return_value = '{"flashcards": []}'
     return lm

tests/test_models.py

@@ -0,0 +1,43 @@
+import json
+
+from app.models import Card, Message
+
+ # Tests for Card and Message models
+def test_card_model():
+    card = Card(question="What is Python?", answer="A programming language")
+    assert card.question == "What is Python?"
+    assert card.answer == "A programming language"
+
+def test_message_model():
+    cards = [
+        Card(question="What is AI?", answer="Artificial Intelligence"),
+        Card(question="What is ML?", answer="Machine Learning")
+    ]
+    message = Message(role="assistant", content=cards)
+    assert message.role == "assistant"
+    assert len(message.content) == 2
+    assert message.content[0].question == "What is AI?"
+
+def test_message_content_json_parsing():
+    json_content = '[{"question": "Q1", "answer": "A1"}, {"question": "Q2", "answer": "A2"}]'
+    message = Message(role="assistant", content=json_content)
+    assert len(message.content) == 2
+    assert message.content[0].question == "Q1"
+
+def test_message_content_to_json():
+    cards = [Card(question="Q1", answer="A1"), Card(question="Q2", answer="A2")]
+    message = Message(role="assistant", content=cards)
+    json_output = message.content_to_json()
+    assert json.loads(json_output) == [
+        {"question": "Q1", "answer": "A1"},
+        {"question": "Q2", "answer": "A2"}
+    ]
+
+# failed test
+def test_message_content_to_csv():
+    cards = [Card(question="Q1", answer="A1"), Card(question="Q2", answer="A2")]
+    message = Message(role="assistant", content=cards)
+    csv_output = message.content_to_csv()
+    expected_output = "Question,Answer\nQ1,A1\nQ2,A2\n"  # Use Unix-style line endings
+    print(csv_output)  # Optional: for debugging purposes
+    assert csv_output == expected_output

tests/test_pipeline.py

@@ -1,18 +1,67 @@
 import pytest
+from unittest.mock import Mock, patch
+import json
+from io import StringIO
+from pydantic import ValidationError
+from app.pipeline import Pipeline, Message, Card, parse_message
+from app.models import PydanticEncoder
 
-def test_generate_flashcards(pipeline, mocker):
-    """
-    Test the generate_flashcards method of LanguageModel.
-    """
-    prompt = "Sample prompt for flashcard generation."
-    expected_response = '{"flashcards": [{"Question": "What is AI?", "Answer": "Artificial Intelligence."}]}'
+# Tests for Pipeline class
+@pytest.fixture
+def mock_pipeline():
+    with patch('app.pipeline') as mock_pipe:
+        mock_pipe.return_value = Mock()
+        yield Pipeline("mock_model")
 
-    # Configure the mock to return a specific response
-    pipeline.generate_flashcards.return_value = expected_response
+# def test_extract_flashcards(mock_pipeline):
+#     mock_pipeline.torch_pipe.return_value = [{"generated_text": [{"role": "assistant", "content": '[{"question": "Q", "answer": "A"}]'}]}]
+#     response = mock_pipeline.extract_flashcards("Test content")
+#     assert isinstance(response, dict)
+#     assert "content" in response
 
-    # Call the method
-    response = pipeline.generate_flashcards(prompt)
+# def test_format_flashcards_csv(mock_pipeline):
+#     response = {"role": "assistant", "content": '[{"question": "Q", "answer": "A"}]'}
+#     formatted = mock_pipeline.format_flashcards("csv", response)
+#     assert formatted.strip() == "Question,Answer\nQ,A"
 
-    # Assertions
-    assert response == expected_response
-    pipeline.generate_flashcards.assert_called_once_with(prompt)
+# def test_generate_flashcards(mock_pipeline):
+#     mock_pipeline.extract_flashcards.return_value = {"role": "assistant", "content": '[{"question": "Q", "answer": "A"}]'}
+#     result = mock_pipeline.generate_flashcards("json", "Test content")
+#     assert json.loads(result) == [{"question": "Q", "answer": "A"}]
+
+# Tests for parse_message function
+def test_parse_message_valid_input():
+    input_dict = {
+        "role": "assistant",
+        "content": '[{"question": "Q1", "answer": "A1"}, {"question": "Q2", "answer": "A2"}]'
+    }
+    message = parse_message(input_dict)
+    assert isinstance(message, Message)
+    assert message.role == "assistant"
+    assert len(message.content) == 2
+
+def test_parse_message_invalid_json():
+    input_dict = {
+        "role": "assistant",
+        "content": 'Invalid JSON'
+    }
+    with pytest.raises(ValueError, match="Invalid JSON in content"):
+        parse_message(input_dict)
+
+def test_parse_message_missing_key():
+    input_dict = {
+        "content": '[{"question": "Q", "answer": "A"}]'
+    }
+    with pytest.raises(ValueError, match="Missing required key"):
+        parse_message(input_dict)
+
+# Test for PydanticEncoder
+def test_pydantic_encoder():
+    card = Card(question="Q", answer="A")
+    encoded = json.dumps(card, cls=PydanticEncoder)
+    assert json.loads(encoded) == {"question": "Q", "answer": "A"}
+
+# Test error cases
+def test_message_invalid_content():
+    with pytest.raises(ValidationError):
+        Message(role="assistant", content="Invalid content")

tests/test_processing.py

@@ -1,71 +1,40 @@
 import pytest
-from app.processing import process_text_input, process_file
-
-def test_process_text_input_success(pipeline):
-    """
-    Test processing of valid text input.
-    """
-    input_text = "This is a sample text for flashcard extraction."
-    output_format = "JSON"
-    expected_output = '{"flashcards": []}'
-
-    result = process_text_input(input_text, output_format, pipeline)
-    assert result == expected_output
-    pipeline.generate_flashcards.assert_called_once()
-
-def test_process_text_input_empty(pipeline):
-    """
-    Test processing of empty text input.
-    """
-    input_text = "   "
-    output_format = "JSON"
-
-    with pytest.raises(ValueError) as excinfo:
-        process_text_input(input_text, output_format, pipeline)
-    assert "No text provided." in str(excinfo.value)
-
-def test_process_file_unsupported_type(pipeline, tmp_path):
-    """
-    Test processing of an unsupported file type.
-    """
-    # Create a dummy unsupported file
-    dummy_file = tmp_path / "dummy.unsupported"
-    dummy_file.write_text("Unsupported content")
-
-    with pytest.raises(ValueError) as excinfo:
-        process_file(dummy_file, "JSON", pipeline)
-    assert "Unsupported file type." in str(excinfo.value)
-
-def test_process_file_pdf(pipeline, tmp_path, mocker):
-    """
-    Test processing of a PDF file.
-    """
-    # Mock the process_pdf function
-    mocker.patch('app.processing.process_pdf', return_value="Extracted PDF text.")
-
-    # Create a dummy PDF file
-    dummy_file = tmp_path / "test.pdf"
-    dummy_file.write_text("PDF content")
-
-    expected_output = '{"flashcards": []}'
-
-    result = process_file(dummy_file, "JSON", pipeline)
-    assert result == expected_output
-    pipeline.generate_flashcards.assert_called_once()
-
-def test_process_file_txt(pipeline, tmp_path, mocker):
-    """
-    Test processing of a TXT file.
-    """
-    # Mock the read_text_file function
-    mocker.patch('app.processing.read_text_file', return_value="Extracted TXT text.")
-
-    # Create a dummy TXT file
-    dummy_file = tmp_path / "test.txt"
-    dummy_file.write_text("TXT content")
-
-    expected_output = '{"flashcards": []}'
-
-    result = process_file(dummy_file, "JSON", pipeline)
-    assert result == expected_output
-    pipeline.generate_flashcards.assert_called_once()
+from unittest.mock import patch, Mock
+from app.processing import process_pdf, read_text_file, process_file, process_text_input
+
+def test_read_text_file_error():
+    with patch("builtins.open", side_effect=IOError("File read error")):
+        with pytest.raises(ValueError, match="Error reading text file: File read error"):
+            read_text_file("test.txt")
+
+# Test for process_file function
+def test_process_file_pdf(pipeline):
+    mock_file = Mock()
+    mock_file.name = "test.pdf"
+
+    with patch('app.processing.process_pdf', return_value="PDF content"):
+        result = process_file(mock_file, "json", pipeline)
+        pipeline.generate_flashcards.assert_called_once_with("json", "PDF content")
+        assert result == '{"flashcards": []}'
+
+def test_process_file_txt(pipeline):
+    mock_file = Mock()
+    mock_file.name = "test.txt"
+
+    with patch('app.processing.read_text_file', return_value="Text content"):
+        result = process_file(mock_file, "json", pipeline)
+        pipeline.generate_flashcards.assert_called_once_with("json", "Text content")
+        assert result == '{"flashcards": []}'
+
+def test_process_file_unsupported():
+    mock_file = Mock()
+    mock_file.name = "test.unsupported"
+
+    with pytest.raises(ValueError, match="Unsupported file type."):
+        process_file(mock_file, "json", None)
+
+# Ensure the pipeline fixture is used in all tests that require it
+@pytest.mark.usefixtures("pipeline")
+class TestWithPipeline:
+    def test_pipeline_usage(self, pipeline):
+        assert pipeline.generate_flashcards.return_value == '{"flashcards": []}'