Update src/app/conversation_core.py

src/app/conversation_core.py

@@ -1,53 +1,43 @@
 
 ###############################################################
-# conversation_core.py — Agentic Partner Core (Qwen 1.5B + Whisper)
+# conversation_core.py — Agentic Partner Core
 ###############################################################
 
 import io
 import re
-import tempfile
 from dataclasses import dataclass
 from typing import List, Optional, Tuple
-from .config import get_user_dir
 
 import numpy as np
-from transformers import pipeline
 from pydub import AudioSegment
-
 import torch
 from gtts import gTTS
 from transformers import (
     AutoTokenizer,
     AutoModelForCausalLM,
+    pipeline,
 )
 
-###############################################################
-# MODEL / LANGUAGE CONSTANTS
-###############################################################
+from .config import get_user_dir
 
-QWEN_MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
-WHISPER_MODEL_SIZE = "base"  # you can change to "large-v3" if GPU budget allows
 
-_QWEN_TOKENIZER = None
-_QWEN_MODEL = None
-_WHISPER = None
+################################################################
+# MODEL CONSTANTS
+################################################################
 
-# Spoken language -> Whisper hint
-WHISPER_LANG_MAP = {
-    "english": "en",
-    "german": "de",
-    "spanish": "es",
-    "russian": "ru",
-    "french": "fr",
-    "italian": "it",
-    "japanese": "ja",
-    "chinese": "zh",
-    "korean": "ko",
-    "arabic": "ar",
-    "hindi": "hi",
+QWEN_MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
+
+# CEFR control hints
+CONTROL_PROMPTS = {
+    "A1": "Use extremely short, simple sentences and very basic vocabulary.",
+    "A2": "Use simple sentences and common everyday vocabulary.",
+    "B1": "Use moderately complex sentences and conversational vocabulary.",
+    "B2": "Use natural, fluent sentences with richer vocabulary.",
+    "C1": "Use complex, advanced sentences with nuanced expressions.",
+    "C2": "Use highly sophisticated, near-native language and style.",
 }
 
-# Spoken language -> gTTS language code
+# spoken language → TTS language
 GTTS_LANG = {
     "english": "en",
     "spanish": "es",
@@ -60,26 +50,24 @@ GTTS_LANG = {
     "italian": "it",
 }
 
-CONTROL_PROMPTS = {
-    "A1": "Use extremely short, simple sentences and very basic vocabulary.",
-    "A2": "Use simple sentences and common everyday vocabulary.",
-    "B1": "Use moderately complex sentences and conversational vocabulary.",
-    "B2": "Use natural, fluent sentences with richer vocabulary.",
-    "C1": "Use complex, advanced sentences with nuanced expressions.",
-    "C2": "Use highly sophisticated, near-native language and style.",
-}
 
+################################################################
+# GLOBAL MODELS
+################################################################
 
-###############################################################
-# GLOBAL LOADERS
-###############################################################
+_QWEN_TOKENIZER = None
+_QWEN_MODEL = None
+_WHISPER_PIPE = None
 
-def load_partner_lm() -> Tuple[AutoTokenizer, AutoModelForCausalLM]:
+
+def load_partner_lm():
+    """Load Qwen conversational model once."""
     global _QWEN_TOKENIZER, _QWEN_MODEL
-    if _QWEN_TOKENIZER is not None and _QWEN_MODEL is not None:
+    if _QWEN_MODEL is not None:
         return _QWEN_TOKENIZER, _QWEN_MODEL
 
-    print("[conversation_core] Loading partner LM:", QWEN_MODEL_NAME)
+    print("[conversation_core] loading:", QWEN_MODEL_NAME)
+
     tok = AutoTokenizer.from_pretrained(QWEN_MODEL_NAME, trust_remote_code=True)
     model = AutoModelForCausalLM.from_pretrained(
         QWEN_MODEL_NAME,
@@ -87,58 +75,30 @@ def load_partner_lm() -> Tuple[AutoTokenizer, AutoModelForCausalLM]:
         device_map="auto",
         trust_remote_code=True,
     )
+
     _QWEN_TOKENIZER = tok
     _QWEN_MODEL = model
     return tok, model
 
 
-# ---------------------------------
-# OPENAI WHISPER (CPU)
-# ---------------------------------
-_whisper_model = None
+def load_whisper_pipe():
+    """Load Whisper ASR pipeline once."""
+    global _WHISPER_PIPE
+    if _WHISPER_PIPE is not None:
+        return _WHISPER_PIPE
 
-_LANG_HINTS = {
-    "english": "en",
-    "spanish": "es",
-    "german": "de",
-    "russian": "ru",
-    "japanese": "ja",
-    "chinese": "zh",
-    "korean": "ko",
-}
-
-
-##########################################
-# SPEECH RECOGNITION — faster-whisper
-##########################################
-
-def transcribe_audio(audio_segment, spoken_lang=None):
-    """
-    Accepts a pydub AudioSegment (mono, 16k).
-    Returns transcript, detected_language, confidence.
-    """
-    global _whisper_model
-    if _whisper_model is None:
-        load_whisper()
-
-    import numpy as np
-
-    audio = np.array(audio_segment.get_array_of_samples()).astype("float32") / 32768.0
-
-    segments, info = _whisper_model.transcribe(
-        audio,
-        beam_size=5,
-        language=spoken_lang,
+    print("[conversation_core] loading Whisper pipeline…")
+    _WHISPER_PIPE = pipeline(
+        "automatic-speech-recognition",
+        model="openai/whisper-small",
+        device="cpu",
     )
-
-    full_text = " ".join([s.text.strip() for s in segments])
-
-    return full_text.strip(), info.language, info.language_probability
+    return _WHISPER_PIPE
 
 
-###############################################################
+################################################################
 # DATA STRUCTURE
-###############################################################
+################################################################
 
 @dataclass
 class ConversationTurn:
@@ -146,104 +106,90 @@ class ConversationTurn:
     text: str
 
 
-###############################################################
-# CLEANING
-###############################################################
+
+################################################################
+# CLEANING LM OUTPUT
+################################################################
 
 def clean_assistant_reply(text: str) -> str:
-    """Strip meta, identity, and obvious junk from LM output."""
+    """Remove meta junk, labels, identity statements."""
     if not text:
         return ""
 
-    # Remove labels
-    text = re.sub(r"(?i)\b(user|assistant|system)\s*:\s*", "", text)
+    # Remove "assistant:" echo
+    text = re.sub(r"(?i)\bassistant\s*:\s*", "", text)
+    text = re.sub(r"(?i)\buser\s*:\s*", "", text)
 
-    # Remove numbered / bullet lists (not wanted in casual chat)
+    # Remove bullet lists (not desired in conversation)
     text = re.sub(r"(?m)^\s*[-•*]\s+.*$", "", text)
     text = re.sub(r"(?m)^\s*\d+\.\s+.*$", "", text)
 
-    # Remove obvious identity / HR / meta nonsense
+    # Remove identity claims
     identity_patterns = [
-        r"(?i)i am (an?|the)? ?(ai|assistant|speaker|model|natural person).*",
-        r"(?i)my name is [A-Za-zäöüÄÖÜß]+.*",
-        r"(?i)i was created.*",
-        r"(?i)human resources manager.*",
-        r"(?i)job description.*",
+        r"(?i)i am an ai.*",
         r"(?i)i am a large language model.*",
+        r"(?i)i was created.*",
+        r"(?i)my name is .*",
     ]
-    for pat in identity_patterns:
-        text = re.sub(pat, "", text)
-
-    # Trim hanging word fragments at the end
-    text = re.sub(r"[A-Za-zÄÖÜäöüß]+$", "", text)
+    for p in identity_patterns:
+        text = re.sub(p, "", text)
 
-    # Collapse whitespace
     text = re.sub(r"\s{2,}", " ", text)
     return text.strip()
 
 
-###############################################################
+
+################################################################
 # CONVERSATION MANAGER
-###############################################################
+################################################################
 
 class ConversationManager:
     def __init__(
         self,
-        target_language: str = "german",
-        native_language: str = "english",
-        cefr_level: str = "B1",
-        topic: str = "general conversation",
+        target_language="german",
+        native_language="english",
+        cefr_level="B1",
+        topic="general conversation",
     ):
-        self.target_language = (target_language or "english").strip().lower()
-        self.native_language = (native_language or "english").strip().lower()
-        self.cefr_level = cefr_level or "B1"
-        self.topic = topic or "general conversation"
+        self.target_language = target_language.lower()
+        self.native_language = native_language.lower()
+        self.cefr_level = cefr_level.upper()
+        self.topic = topic
         self.history: List[ConversationTurn] = []
 
-        # Warm-load models once per session
         load_partner_lm()
-        load_whisper()
+        load_whisper_pipe()
 
-    ###########################################################
-    # PROMPT + GENERATION
-    ###########################################################
+    ################################################################
+    # SYSTEM PROMPT
+    ################################################################
 
-    def _build_system_prompt(self) -> str:
+    def _build_system_prompt(self):
         base = (
             f"You are a friendly conversation partner speaking {self.target_language}. "
             f"Reply ONLY in {self.target_language}. "
-            f"Do NOT explain grammar, vocabulary, or translations unless the user explicitly asks. "
-            f"Do NOT describe what the sentence means, do NOT say 'the sentence translates to...', "
-            f"and do NOT mention that you are explaining anything. "
             f"Adapt your language to CEFR level {self.cefr_level}. "
             f"{CONTROL_PROMPTS.get(self.cefr_level, '')} "
-            "Keep your replies natural and conversational, usually 1–3 short sentences. "
-            "Ask exactly ONE natural follow-up question related to what the user said. "
-            "Never end the conversation unless the user explicitly ends it. "
-            "Do NOT say goodbye or conclude unless the user does. "
-            "Never talk about being an AI, model, or assistant. "
-            "Do not mention job descriptions, resumes, or HR responsibilities unless the user clearly asks. "
+            f"Topic of conversation: {self.topic}. "
+            "Give 1–3 short natural sentences and ALWAYS end with 1 follow-up question. "
+            "Never mention AI, assistants, grammar explanations, or meta commentary."
         )
-        if self.topic.strip():
-            base += f"The main topic of conversation is: {self.topic.strip()}. "
         return base
 
+    ################################################################
+    # GENERATION
+    ################################################################
+
     def _generate_lm(self, user_text: str) -> str:
         tok, model = load_partner_lm()
 
-        system_prompt = self._build_system_prompt()
         messages = [
-            {"role": "system", "content": system_prompt},
-            {
-                "role": "user",
-                "content": f"The user (who speaks {self.native_language}) said: {user_text}",
-            },
+            {"role": "system", "content": self._build_system_prompt()},
+            {"role": "user", "content": user_text},
         ]
 
         prompt = tok.apply_chat_template(
-            messages,
-            tokenize=False,
-            add_generation_prompt=True,
+            messages, tokenize=False, add_generation_prompt=True
         )
 
         enc = tok(prompt, return_tensors="pt").to(model.device)
@@ -251,66 +197,49 @@ class ConversationManager:
         with torch.no_grad():
             out = model.generate(
                 **enc,
-                max_new_tokens=160,           # enough space for natural replies
+                max_new_tokens=160,
                 temperature=0.8,
                 top_p=0.95,
-                top_k=50,
                 repetition_penalty=1.15,
-                pad_token_id=tok.eos_token_id,
                 do_sample=True,
+                pad_token_id=tok.eos_token_id,
             )
 
-        raw = tok.decode(out[0], skip_special_tokens=True).strip()
-
-        # If the user text is echoed, strip it
-        if user_text in raw:
-            raw = raw.split(user_text)[-1].strip()
-
-        # Remove "assistant" label echoes
-        lines = [
-            ln for ln in raw.splitlines()
-            if ln.strip().lower() not in ("assistant", "assistant:")
-        ]
-        raw = "\n".join(lines).strip()
+        raw = tok.decode(out[0], skip_special_tokens=True)
 
-        return clean_assistant_reply(raw)
+        # Remove echo
+        cleaned = clean_assistant_reply(raw)
+        return cleaned
 
-    ###########################################################
+    ################################################################
     # PUBLIC REPLY API
-    ###########################################################
+    ################################################################
 
-    def reply(self, user_text: str, input_lang: str = "german"):
-        """Generate a reply + explanation + TTS audio."""
+    def reply(self, user_text: str, input_lang="german"):
         self.history.append(ConversationTurn("user", user_text))
 
         assistant_text = self._generate_lm(user_text)
         self.history.append(ConversationTurn("assistant", assistant_text))
 
         explanation = self._generate_explanation(assistant_text)
-        audio = self.text_to_speech(assistant_text)
+        audio_bytes = self.text_to_speech(assistant_text)
 
         return {
             "reply_text": assistant_text,
             "explanation": explanation,
-            "audio": audio,
+            "audio": audio_bytes,
         }
 
-    ###########################################################
-    # SHORT EXPLANATION (EN / native language)
-    ###########################################################
+    ################################################################
+    # SHORT EXPLANATION
+    ################################################################
 
     def _generate_explanation(self, assistant_text: str) -> str:
-        """Return exactly ONE simple native-language sentence, no meta, no logic."""
-        if not assistant_text:
-            return ""
-
         tok, model = load_partner_lm()
+
         prompt = (
             f"Rewrite the meaning of this {self.target_language} sentence "
-            f"in ONE very short {self.native_language} sentence. "
-            f"Do NOT explain what you are doing, do NOT say 'the sentence means', "
-            f"do NOT describe tone, and do NOT provide multiple versions.\n"
-            f"Sentence: \"{assistant_text}\""
+            f"in ONE short {self.native_language} sentence:\n{assistant_text}"
         )
 
         enc = tok(prompt, return_tensors="pt").to(model.device)
@@ -323,60 +252,39 @@ class ConversationManager:
                 pad_token_id=tok.eos_token_id,
             )
 
-        raw = tok.decode(out[0], skip_special_tokens=True)
-        raw = raw.replace(prompt, "").strip()
-
-        # keep first sentence only
-        parts = re.split(r"(?<=[.!?])\s+", raw)
-        if parts:
-            raw = parts[0].strip()
+        decoded = tok.decode(out[0], skip_special_tokens=True)
+        cleaned = decoded.replace(prompt, "").strip()
 
-        # remove meta leftovers
-        raw = re.sub(r"(?i)the sentence.*$", "", raw)
-        raw = re.sub(r"(?i)this means.*$", "", raw)
+        # keep only the first sentence
+        parts = re.split(r"(?<=[.!?])\s+", cleaned)
+        return parts[0].strip()
 
-        return raw.strip()
+    ################################################################
+    # TRANSCRIPTION — SINGLE VALID VERSION
+    ################################################################
 
-
-
-    ###########################################################
-    # AUDIO TRANSCRIPTION — Transformers Whisper
-    ###########################################################
-    
-    from transformers import pipeline
-    
-    ###########################################################
-    # AUDIO TRANSCRIPTION — Transformers Whisper
-    ###########################################################
-    
-    whisper_pipe = pipeline(
-        task="automatic-speech-recognition",
-        model="openai/whisper-small",
-        device="cpu"
-    )
-    
     def transcribe(self, audio_segment, spoken_lang=None):
-        import numpy as np
-    
+        """Transcribe using Transformers Whisper."""
+        pipe = load_whisper_pipe()
+
         audio = np.array(audio_segment.get_array_of_samples()).astype("float32")
-        audio = audio / np.max(np.abs(audio))
-    
-        result = whisper_pipe(audio)
+        audio = audio / max(np.max(np.abs(audio)), 1e-6)
+
+        result = pipe(audio)
         text = result.get("text", "").strip()
-    
+
         return text, spoken_lang or "unknown", 1.0
 
-    ###########################################################
-    # TEXT → SPEECH
-    ###########################################################
+    ################################################################
+    # TTS — gTTS
+    ################################################################
 
     def text_to_speech(self, text: str) -> Optional[bytes]:
-        """Return MP3 bytes for the assistant text, or None on failure."""
         if not text:
             return None
         try:
-            lang_code = GTTS_LANG.get(self.target_language, "en")
-            tts = gTTS(text=text, lang=lang_code)
+            lang = GTTS_LANG.get(self.target_language, "en")
+            tts = gTTS(text=text, lang=lang)
             buf = io.BytesIO()
             tts.write_to_fp(buf)
             return buf.getvalue()
@@ -384,8 +292,6 @@ class ConversationManager:
             return None
 
 
-###############################################################
+################################################################
 # END OF FILE
-###############################################################
-
-
+################################################################