Update app.py

app.py

@@ -6,11 +6,10 @@ import torch
 import uvicorn
 import os
 
-# Set cache directories to /tmp which is writable
 os.environ['TRANSFORMERS_CACHE'] = '/tmp/transformers_cache'
 os.environ['TORCH_HOME'] = '/tmp/torch_cache'
 
-app = FastAPI(title="TinyLlama Fitness Bot")
+app = FastAPI(title="DIANA - Diet And Nutrition Assistant")
 
 app.add_middleware(
     CORSMiddleware,
@@ -20,7 +19,6 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Global variables
 model = None
 tokenizer = None
 MODEL_LOADED = False
@@ -30,81 +28,120 @@ def load_model():
     try:
         print("Starting model load...")
         model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-        
-        # CPU-specific settings
         torch.set_num_threads(4)
         
-        print("Loading tokenizer...")
         tokenizer = AutoTokenizer.from_pretrained(
             model_name,
-            cache_dir='/tmp/transformers_cache'  # Use /tmp directory
+            cache_dir='/tmp/transformers_cache'
         )
         
-        print("Loading model...")
         model = AutoModelForCausalLM.from_pretrained(
             model_name,
             torch_dtype=torch.float32,
             low_cpu_mem_usage=True,
-            device_map=None,  # Force CPU
-            cache_dir='/tmp/transformers_cache'  # Use /tmp directory
+            device_map=None,
+            cache_dir='/tmp/transformers_cache'
         )
         
         model.eval()
         MODEL_LOADED = True
-        print("Model loaded successfully on CPU!")
         return True
     except Exception as e:
         print(f"Error loading model: {str(e)}")
         MODEL_LOADED = False
         return False
 
-# Load model on startup
-print("Initiating model load...")
+print("Initiating DIANA...")
 load_model()
 
 class Query(BaseModel):
     prompt: str
-    max_length: int = 100
+    max_length: int = 200
     temperature: float = 0.7
 
+def is_greeting(text):
+    greetings = ['hi', 'hello', 'hey', 'good morning', 'good afternoon', 'good evening', 'greetings']
+    return any(greeting in text.lower() for greeting in greetings)
+
+def is_fitness_question(text):
+    fitness_keywords = [
+        'workout', 'exercise', 'training', 'muscle', 'strength', 'cardio', 'weight',
+        'diet', 'nutrition', 'protein', 'carbs', 'fat', 'meal', 'food', 'eating',
+        'routine', 'program', 'sets', 'reps', 'gym', 'fitness', 'health'
+    ]
+    return any(keyword in text.lower() for keyword in fitness_keywords)
+
 @app.post("/chat")
 async def chat(query: Query):
-    global model, tokenizer, MODEL_LOADED
-    
     if not MODEL_LOADED:
         if not load_model():
             raise HTTPException(
                 status_code=503, 
-                detail="Model is not loaded. Please try again in a minute."
+                detail="DIANA is still initializing. Please try again in a minute."
             )
     
     try:
-        formatted_prompt = f"<|user|>{query.prompt}</s><|assistant|>"
+        # Personalized system prompts
+        if is_greeting(query.prompt):
+            system_prompt = """You are DIANA (Diet And Nutrition Assistant), a friendly and knowledgeable 
+            fitness companion. Always respond warmly and offer to help with fitness and nutrition guidance. 
+            Sign your responses with '- DIANA 💪'"""
+        else:
+            system_prompt = """You are DIANA (Diet And Nutrition Assistant), a knowledgeable fitness and 
+            nutrition guide. Provide practical, safe, and evidence-based advice about workouts, nutrition, 
+            and healthy living. Include:
+            1. Clear, actionable recommendations
+            2. Safety considerations
+            3. Beginner-friendly explanations
+            Remember to sign your responses with '- DIANA 💪'"""
+        
+        formatted_prompt = f"""<|system|>{system_prompt}</s>
+<|user|>{query.prompt}</s>
+<|assistant|>"""
         
         inputs = tokenizer(
             formatted_prompt,
             return_tensors="pt",
             truncation=True,
-            max_length=256
+            max_length=300
         )
         
         with torch.no_grad():
             outputs = model.generate(
                 inputs["input_ids"],
-                max_new_tokens=query.max_length,
-                temperature=query.temperature,
+                max_new_tokens=200,
+                min_new_tokens=50,
+                temperature=0.7,
                 top_p=0.9,
                 do_sample=True,
                 pad_token_id=tokenizer.eos_token_id,
-                num_beams=1,
+                repetition_penalty=1.2,
+                no_repeat_ngram_size=3,
+                eos_token_id=tokenizer.eos_token_id,
                 early_stopping=True
             )
         
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
         response = response.split("<|assistant|>")[-1].strip()
         
-        if not response or len(response.split()) < 3:
-            return {"response": "I apologize, could you please rephrase your question?"}
+        # Add signature if not present
+        if "- DIANA 💪" not in response:
+            response = response + "\n\n- DIANA 💪"
+        
+        # Response validation and fallbacks
+        if not response or len(response.split()) < 20:
+            if is_greeting(query.prompt):
+                return {
+                    "response": "Hi there! I'm DIANA, your personal Diet And Nutrition Assistant. I'm here to help you achieve your health and fitness goals! Would you like some advice about workouts or nutrition?\n\n- DIANA 💪"
+                }
+            elif is_fitness_question(query.prompt):
+                return {
+                    "response": "Let me help you on your fitness journey! Could you provide more details about your specific goals and current fitness level? This will help me give you the most relevant advice.\n\n- DIANA 💪"
+                }
+            else:
+                return {
+                    "response": "Hi! I'm DIANA, your Diet And Nutrition Assistant. I specialize in workout plans, diet advice, and general health tips. What would you like to know more about?\n\n- DIANA 💪"
+                }
             
         return {"response": response}
         
@@ -115,18 +152,9 @@ async def chat(query: Query):
 @app.get("/")
 def read_root():
     return {
-        "status": "API is running!",
-        "model_loaded": MODEL_LOADED,
-        "backend": "CPU"
-    }
-
-@app.get("/debug")
-def debug_info():
-    return {
+        "status": "DIANA (Diet And Nutrition Assistant) is running!",
         "model_loaded": MODEL_LOADED,
-        "device": "cpu",
-        "num_threads": torch.get_num_threads(),
-        "cache_dir": os.environ.get('TRANSFORMERS_CACHE')
+        "specialties": ["Personalized workout advice", "Nutrition guidance", "Fitness planning"]
     }
 
 if __name__ == "__main__":