Upload 19 files

src/writing_studio/core/analyzer.py

@@ -67,16 +67,30 @@ class WritingAnalyzer:
                 logger.info(f"Loading new model: {model_name}")
                 self.model_service.load_model(model_name)
 
-            # Generate prompt
+            # Generate prompt using selected pack
             prompt = self.prompt_service.generate_prompt(user_text, prompt_pack)
 
-            # Generate revision
+            # Generate AI revision
+            logger.info("Generating AI revision...")
             with generation_duration.time():
-                revision = self.model_service.generate_text(prompt)
-
-            # Extract only the revised part (after "Revised Text:")
-            if "Revised Text:" in revision:
-                revision = revision.split("Revised Text:")[-1].strip()
+                revision = self.model_service.generate_text(
+                    prompt,
+                    max_length=min(len(user_text.split()) * 2 + 100, settings.max_model_length),
+                    use_cache=True
+                )
+
+            # Clean up revision (remove any prompt artifacts)
+            if prompt_pack in revision:
+                revision = revision.split(prompt_pack)[-1].strip()
+            if "Revised text:" in revision:
+                revision = revision.split("Revised text:")[-1].strip()
+            if user_text in revision:
+                # Model might include original text, extract just the revision
+                revision = revision.replace(user_text, "").strip()
+
+            # If revision is empty or too similar, provide a note
+            if not revision or revision == user_text:
+                revision = user_text + "\n\n[Note: The AI model kept the text as-is, suggesting it's already well-written!]"
 
             # Analyze with rubric
             rubric_results = self.rubric_service.analyze_text(user_text)

src/writing_studio/core/config.py

@@ -42,9 +42,12 @@ class Settings(BaseSettings):
     server_workers: int = Field(default=4, ge=1, description="Number of worker processes")
 
     # Model Configuration
-    default_model: str = Field(default="distilgpt2", description="Default HuggingFace model")
+    default_model: str = Field(
+        default="google/flan-t5-base",
+        description="Default HuggingFace model (instruction-tuned for revision)"
+    )
     max_model_length: int = Field(default=512, ge=1, description="Maximum model input length")
-    default_max_length: int = Field(default=300, ge=1, description="Default generation length")
+    default_max_length: int = Field(default=512, ge=1, description="Default generation length")
     default_num_sequences: int = Field(default=1, ge=1, description="Number of sequences")
 
     # Security

src/writing_studio/services/model_service.py

@@ -5,7 +5,7 @@ import time
 from functools import lru_cache
 from typing import Any, Dict, Optional
 
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 
 from writing_studio.core.config import settings
 from writing_studio.core.exceptions import ModelLoadError, TextGenerationError
@@ -20,6 +20,7 @@ class ModelService:
         """Initialize the model service."""
         self._current_model: Optional[Any] = None
         self._current_model_name: Optional[str] = None
+        self._task_type: str = "text2text-generation"  # Default for FLAN-T5
         self._cache: Dict[str, Any] = {}
         self._load_default_model()
 
@@ -57,13 +58,24 @@ class ModelService:
             logger.info(f"Loading model: {model_name}")
             start_time = time.time()
 
+            # Detect model type and use appropriate pipeline
+            # FLAN-T5, T5 = text2text-generation
+            # GPT-2, GPT = text-generation
+            if any(x in model_name.lower() for x in ['t5', 'flan']):
+                task = "text2text-generation"
+                logger.info(f"Detected instruction-following model, using {task} pipeline")
+            else:
+                task = "text-generation"
+                logger.info(f"Detected text generation model, using {task} pipeline")
+
             # Load model with error handling
-            # Note: cache_dir is handled automatically by transformers
             self._current_model = pipeline(
-                "text-generation",
+                task,
                 model=model_name,
+                max_length=settings.max_model_length,
             )
             self._current_model_name = model_name
+            self._task_type = task
 
             load_time = time.time() - start_time
             logger.info(f"Model loaded successfully in {load_time:.2f}s: {model_name}")
@@ -119,19 +131,32 @@ class ModelService:
             logger.info(f"Generating text with model: {self._current_model_name}")
             start_time = time.time()
 
-            # Generate text with proper parameters
-            result = self._current_model(
-                prompt,
-                max_length=params["max_length"],
-                num_return_sequences=params["num_sequences"],
-                do_sample=True,
-                temperature=params["temperature"],
-                pad_token_id=self._current_model.tokenizer.eos_token_id,  # Avoid warnings
-            )
+            # Generate text with parameters appropriate for model type
+            if self._task_type == "text2text-generation":
+                # T5/FLAN-T5 models
+                result = self._current_model(
+                    prompt,
+                    max_new_tokens=params["max_length"],
+                    num_return_sequences=params["num_sequences"],
+                    do_sample=True,
+                    temperature=params["temperature"],
+                    truncation=True,
+                )
+                # T5 models return generated_text directly
+                generated_text = result[0]["generated_text"]
+            else:
+                # GPT-2 style models
+                result = self._current_model(
+                    prompt,
+                    max_length=params["max_length"],
+                    num_return_sequences=params["num_sequences"],
+                    do_sample=True,
+                    temperature=params["temperature"],
+                    pad_token_id=self._current_model.tokenizer.eos_token_id,
+                )
+                generated_text = result[0]["generated_text"]
 
-            generated_text = result[0]["generated_text"]
             generation_time = time.time() - start_time
-
             logger.info(f"Text generated in {generation_time:.2f}s")
 
             # Cache result if enabled

src/writing_studio/services/prompt_service.py

@@ -9,27 +9,28 @@ class PromptService:
     """Service for managing and generating prompts."""
 
     def __init__(self):
-        """Initialize the prompt service with templates."""
+        """Initialize the prompt service with templates for instruction-following models."""
+        # Optimized for FLAN-T5 and other instruction-tuned models
         self.prompt_packs = {
             "General": {
-                "instruction": "Revise this text for clarity, conciseness, and audience fit",
-                "context": "Focus on improving overall readability and effectiveness.",
+                "instruction": "Revise the following text to improve clarity, conciseness, and readability",
+                "context": "Make it clear and easy to understand while maintaining the original meaning.",
             },
             "Literature": {
-                "instruction": "Revise this literary analysis with attention to theme, style, and evidence",
-                "context": "Ensure proper use of literary terminology and textual support.",
+                "instruction": "Revise this literary analysis to strengthen the argument with better evidence and literary terminology",
+                "context": "Enhance academic rigor and use of textual support.",
             },
             "Tech Comm": {
                 "instruction": "Revise this technical document for precision, clarity, and professional tone",
-                "context": "Emphasize accuracy, clear instructions, and appropriate technical level.",
+                "context": "Make it accurate, clear, and appropriate for technical audiences.",
             },
             "Academic": {
-                "instruction": "Revise this academic writing for formal tone, organization, and scholarly support",
-                "context": "Maintain formal register and ensure proper citation indicators.",
+                "instruction": "Revise this academic writing to improve formal tone, organization, and scholarly voice",
+                "context": "Ensure formal register and proper academic style.",
             },
             "Creative": {
-                "instruction": "Revise this creative writing with focus on imagery, voice, and engagement",
-                "context": "Enhance descriptive language and narrative flow.",
+                "instruction": "Revise this creative writing to enhance imagery, voice, and reader engagement",
+                "context": "Improve descriptive language and narrative flow.",
             },
         }
 
@@ -46,6 +47,8 @@ class PromptService:
         """
         Generate a complete prompt from user text and pack template.
 
+        Optimized for instruction-following models like FLAN-T5.
+
         Args:
             user_text: User's input text
             pack_name: Name of the prompt pack to use
@@ -60,14 +63,8 @@ class PromptService:
         pack = self.prompt_packs[pack_name]
         logger.info(f"Generating prompt with pack: {pack_name}")
 
-        prompt = f"""{pack['instruction']}.
-
-Context: {pack['context']}
-
-Original Text:
-{user_text}
-
-Revised Text:"""
+        # Format optimized for FLAN-T5 and similar instruction-tuned models
+        prompt = f"{pack['instruction']}. {pack['context']}\n\nText: {user_text}\n\nRevised text:"
 
         return prompt