test_llm_allergens.py 2.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879
  1. import os
  2. import sys
  3. import paramiko
  4. import dotenv
  5. repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
  6. dotenv.load_dotenv(dotenv_path=os.path.join(repo_root, ".env"))
  7. host = os.environ.get('SERVER_HOST')
  8. user = os.environ.get('SERVER_USER')
  9. password = os.environ.get('SERVER_PASS')
  10. if password == "your_db_password_here" or not password:
  11. password = None
  12. ssh = paramiko.SSHClient()
  13. ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
  14. try:
  15. ssh.connect(host, username=user, password=password, timeout=10)
  16. print("SSH Connected!")
  17. # Run a test inside the container
  18. python_cmd = """
  19. import ollama
  20. model = 'llama3.2:3b'
  21. # Test product: Peanut butter with milk
  22. name = "Snickers Bar"
  23. ingredients = "milk chocolate, peanuts, sugar, milk, cocoa butter, lactose"
  24. allergens = [
  25. "Peanuts", "Eggs", "Milk", "Wheat", "Gluten", "Soy",
  26. "Tree Nuts", "Fish", "Shellfish", "Sesame", "Mustard",
  27. "Celery", "Lupin", "Molluscs", "Sulphites"
  28. ]
  29. print("--- Testing Loop ---")
  30. import time
  31. t0 = time.time()
  32. detected = []
  33. # Just test first 3 allergens in loop to measure time
  34. for allergen in allergens[:3]:
  35. prompt = f'Answer by yes or no, if it is in some case answer yes : Are {allergen} allergens in the product "{name}" with ingredients "{ingredients}"?'
  36. res = ollama.chat(model=model, messages=[{'role': 'user', 'content': prompt}])
  37. ans = res['message']['content'].strip()
  38. print(f"{allergen}: {ans}")
  39. if 'yes' in ans.lower():
  40. detected.append(allergen)
  41. print("Loop time for 3:", time.time() - t0)
  42. print("\\n--- Testing Single Call ---")
  43. t0 = time.time()
  44. prompt = f'''Given the product "{name}" with ingredients "{ingredients}".
  45. For each allergen in the list: Peanuts, Eggs, Milk, Wheat, Gluten, Soy, Tree Nuts, Fish, Shellfish, Sesame, Mustard, Celery, Lupin, Molluscs, Sulphites.
  46. Answer by yes or no, if it is in some case answer yes: Are they allergens in this product?
  47. Format your response exactly as:
  48. AllergenName: Yes/No
  49. '''
  50. res = ollama.chat(model=model, messages=[{'role': 'user', 'content': prompt}])
  51. print(res['message']['content'].strip())
  52. print("Single call time:", time.time() - t0)
  53. """
  54. # Write python_cmd to a temporary file on the server and copy to container
  55. sftp = ssh.open_sftp()
  56. with sftp.file('food_project/test_ollama.py', 'w') as f:
  57. f.write(python_cmd)
  58. sftp.close()
  59. # Copy file into container
  60. ssh.exec_command("docker cp food_project/test_ollama.py food_project-app-1:/app/test_ollama.py")
  61. stdin, stdout, stderr = ssh.exec_command("docker exec food_project-app-1 python /app/test_ollama.py")
  62. print("STDOUT:\n", stdout.read().decode('utf-8'))
  63. print("STDERR:\n", stderr.read().decode('utf-8'))
  64. except Exception as e:
  65. print(f"Error: {e}")
  66. finally:
  67. ssh.close()