feat: flexible modal message handling

.gitignore

@@ -1 +1,2 @@
 __pycache__
+venv

ai_engine.py

@@ -1,44 +1,165 @@
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+from typing import List, Tuple
 
-print("Loading model...")
+# Use a very small model for testing
+class UniversalChatModel:
+    def __init__(self, model_name: str):
+        self.model_name = model_name
+        print(f"Loading tokenizer for {model_name}...")
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        
+        # Set padding token
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token or "<|endoftext|>"
+            
+        print(f"Loading model {model_name}...")
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float16,
+            device_map="auto"
+        )
+        
+        print("Model loaded successfully!")
+        
+    def format_prompt_fallback(self, messages: List[dict]) -> str:
+        """Universal ChatML format for models without chat templates"""
+        chatml = ""
+        for message in messages:
+            role = message["role"]
+            content = message["content"]
+            chatml += f"<|im_start|>{role}\n{content}<|im_end|>\n"
+        chatml += "<|im_start|>assistant\n"
+        return chatml
+        
+    def build_messages(self, history: List[Tuple[str, str]], current_message: str, system_prompt: str = None) -> List[dict]:
+        """Build universal message format"""
+        messages = []
+        
+        # Add system prompt
+        if system_prompt:
+            messages.append({"role": "system", "content": system_prompt})
+        
+        # Add history
+        for user_msg, assistant_msg in history:
+            messages.append({"role": "user", "content": user_msg})
+            messages.append({"role": "assistant", "content": assistant_msg})
+        
+        # Add current message
+        messages.append({"role": "user", "content": current_message})
+        
+        return messages
+        
+    def format_prompt(self, messages: List[dict]) -> str:
+        """Format prompt using model's chat template or fallback"""
+        # Try model's built-in chat template
+        if hasattr(self.tokenizer, 'chat_template') and self.tokenizer.chat_template:
+            try:
+                prompt = self.tokenizer.apply_chat_template(
+                    messages, 
+                    tokenize=False, 
+                    add_generation_prompt=True
+                )
+                return prompt
+            except Exception as e:
+                print(f"Warning: Failed to use chat template: {e}")
+                pass
+        
+        # Fallback to ChatML
+        return self.format_prompt_fallback(messages)
+        
+    def extract_response(self, prompt: str, generated_text: str) -> str:
+        """Universal response extraction"""
+        if hasattr(self.tokenizer, 'chat_template') and self.tokenizer.chat_template:
+            # Model has chat template - extract after prompt
+            if generated_text.startswith(prompt):
+                response = generated_text[len(prompt):].strip()
+                # Remove any remaining template markers
+                response = response.replace("<|im_end|>", "").strip()
+                response = response.replace("</s>", "").strip()
+                return response
+            else:
+                # Fallback: find first assistant response
+                if "assistant" in generated_text.lower():
+                    parts = generated_text.lower().split("assistant")
+                    response = parts[-1].strip()
+                    response = response.replace("<|im_end|>", "").strip()
+                    response = response.replace("</s>", "").strip()
+                    return response
+                return generated_text.strip()
+        else:
+            # ChatML fallback
+            if "<|im_start|>assistant\n" in generated_text:
+                parts = generated_text.split("<|im_start|>assistant\n")
+                response = parts[-1].replace("<|im_end|>", "").strip()
+                response = response.replace("</s>", "").strip()
+                return response
+            elif generated_text.startswith(prompt):
+                response = generated_text[len(prompt):].strip()
+                response = response.replace("<|im_end|>", "").strip()
+                response = response.replace("</s>", "").strip()
+                return response
+            else:
+                return generated_text.strip()
+        
+    def generate(self, message: str, history: List[Tuple[str, str]] = None, system_prompt: str = None) -> str:
+        """Generate response using universal chat template system"""
+        if history is None:
+            history = []
+            
+        # Build messages
+        messages = self.build_messages(history, message, system_prompt)
+        
+        # Format prompt
+        prompt = self.format_prompt(messages)
+        
+        print(f"\n\n----- PROMPT -----\n{prompt}\n-----------------\n\n")
+        
+        # Tokenize
+        inputs = self.tokenizer(prompt, return_tensors="pt", padding=True)
+        # Move to model device
+        inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
+        
+        # Generate
+        generation_config = {
+            "max_new_tokens": 150,
+            "do_sample": True,
+            "temperature": 0.7,
+            "pad_token_id": self.tokenizer.eos_token_id,
+            "eos_token_id": self.tokenizer.eos_token_id,
+        }
+        
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                **generation_config
+            )
+        
+        # Decode
+        generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        print(f"\n\n----- GENERATED -----\n{generated_text}\n-------------------\n\n")
+        
+        # Extract response
+        response = self.extract_response(prompt, generated_text)
+        
+        return response
 
-generator = pipeline(
-    "text-generation",
-    model="Qwen/Qwen2.5-0.5B-Instruct"
-    # model="google/flan-t5-small"
-)
+# Initialize with a tiny model for testing
+# MODEL_NAME = "HuggingFaceH4/tiny-random-LlamaForCausalLM"
+MODEL_NAME = "google/functiongemma-270m-it"
+SYSTEM_PROMPT = "You are a helpful assistant."
 
-SYSTEM_PROMPT = (
-    "You are a helpful assistant. "
-    "Answer clearly and concisely."
-)
+# Create global model instance
+print("Creating model instance...")
+chat_model = UniversalChatModel(MODEL_NAME)
 
-def build_prompt(history: list[str], message: str) -> str:
-    prompt = SYSTEM_PROMPT + "\n\n"
+def generate(message: str, history: List[Tuple[str, str]]) -> str:
+    """Generate response using universal chat model"""
+    return chat_model.generate(message, history, SYSTEM_PROMPT)
 
-    for line in history:
-        prompt += line + "\n"
-
-    prompt += f"User: {message}\nAI:"
-    return prompt
-
-def generate(message: str, history: list[str]) -> str:
-    prompt = build_prompt(history, message)
-
-    result = generator(
-        prompt,
-        max_new_tokens=500,
-        do_sample=True
-        # temperature=0.7
-    )
-
-    text = result[0]["generated_text"]
-
-    # Extract only the newly generated text (after the prompt)
-    if text.startswith(prompt):
-        response = text[len(prompt):].strip()
-    else:
-        # Fallback to original logic with better error handling
-        response = text.split("AI:")[-1].strip()
-    
-    return response
+if __name__ == "__main__":
+    # Quick test
+    print("Testing generation...")
+    reply = generate("What is 2+2?", [])
+    print(f"Final response: {reply}")

api.py

@@ -4,8 +4,8 @@ from ai_engine import generate
 
 app = FastAPI(title="Hello World AI")
 
-# VERY simple in-memory store
-sessions: dict[str, list[str]] = {}
+# VERY simple in-memory store - now storing tuples (user_msg, assistant_msg)
+sessions: dict[str, list[tuple[str, str]]] = {}
 
 
 class ChatRequest(BaseModel):
@@ -23,7 +23,7 @@ def chat(req: ChatRequest):
 
     reply = generate(req.message, history)
 
-    history.append(f"User: {req.message}")
-    history.append(f"AI: {reply}")
+    # Store as tuple (user_msg, assistant_msg)
+    history.append((req.message, reply))
 
     return {"reply": reply}

requirements.txt

@@ -1,3 +1,4 @@
+accelerate
 fastapi
 uvicorn
 transformers