
from __future__ import annotations

import math
from collections import Counter
from string import ascii_lowercase




def    __UpperCamelCase     (				lowerCAmelCase__							:				str  ):
 __a				,       __a :  int						=  analyze_text(lowerCAmelCase__  )
 __a :  Union[str, Any]						=  list(''' ''' + ascii_lowercase  )
 # what is our total sum of probabilities.
 __a :  Union[str, Any]						=  sum(single_char_strings.values()  )

 # one length string
 __a :  int						=  0
 # for each alpha we go in our dict and if it is in it we calculate entropy
 for ch in my_alphas:
  if ch in single_char_strings:
   __a :  Any						=  single_char_strings[ch]
   __a :  List[Any]						=  my_str / all_sum
   my_fir_sum += prob * math.loga(lowerCAmelCase__  )  # entropy formula.

    # print entropy
 print(f"{round(-1 * my_fir_sum  ):.1f}"  )

 # two len string
 __a :  Tuple						=  sum(two_char_strings.values()  )
 __a :  Optional[int]						=  0
 # for each alpha (two in size) calculate entropy.
 for cha in my_alphas:
  for cha in my_alphas:
   __a :  List[Any]						=  cha + cha
   if sequence in two_char_strings:
    __a :  List[Any]						=  two_char_strings[sequence]
    __a :  int						=  int(lowerCAmelCase__  ) / all_sum
    my_sec_sum += prob * math.loga(lowerCAmelCase__  )

    # print second entropy
 print(f"{round(-1 * my_sec_sum  ):.1f}"  )

 # print the difference between them
 print(f"{round((-1 * my_sec_sum) - (-1 * my_fir_sum)  ):.1f}"  )




def    __UpperCamelCase     (				lowerCAmelCase__							:				str  ):
 __a :  Any						=  Counter()  # type: ignore
 __a :  int						=  Counter()  # type: ignore
 single_char_strings[text[-1]] += 1

 # first case when we have space at start.
 two_char_strings[" " + text[0]] += 1
 for i in range(0  ,   len(lowerCAmelCase__  ) - 1  ):
  single_char_strings[text[i]] += 1
  two_char_strings[text[i : i + 2]] += 1
 return single_char_strings, two_char_strings




def    __UpperCamelCase     (				):
 import doctest

 doctest.testmod()
 # text = (
 #     "Had repulsive dashwoods suspicion sincerity but advantage now him. Remark "
 #     "easily garret nor nay. Civil those mrs enjoy shy fat merry. You greatest "
 #     "jointure saw horrible. He private he on be imagine suppose. Fertile "
 #     "beloved evident through no service elderly is. Blind there if every no so "
 #     "at. Own neglected you preferred way sincerity delivered his attempted. To "
 #     "of message cottage windows do besides against uncivil.  Delightful "
 #     "unreserved impossible few estimating men favourable see entreaties. She "
 #     "propriety immediate was improving. He or entrance humoured likewise "
 #     "moderate. Much nor game son say feel. Fat make met can must form into "
 #     "gate. Me we offending prevailed discovery. "
 # )

 # calculate_prob(text)


if __name__ == "__main__":
      main()


import gc
import unittest

from diffusers import FlaxDPMSolverMultistepScheduler, FlaxStableDiffusionPipeline
from diffusers.utils import is_flax_available, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
      import jax
      import jax.numpy as jnp
      from flax.jax_utils import replicate
      from flax.training.common_utils import shard



@slow
@require_flax
class 		UpperCamelCase__  (       unittest.TestCase    ):
   def       lowerCAmelCase			(self :					Optional[Any]   ):
    # clean up the VRAM after each test
    super().tearDown()
    gc.collect()
   def       lowerCAmelCase			(self :					int   ):
    __a				,       __a :  Optional[Any]						=  FlaxStableDiffusionPipeline.from_pretrained(
        '''stabilityai/stable-diffusion-2'''					,				revision='''bf16'''					,				dtype=jnp.bfloataa					,				)

    __a :  Any						=  '''A painting of a squirrel eating a burger'''
    __a :  Dict						=  jax.device_count()
    __a :  Optional[int]						=  num_samples * [prompt]
    __a :  Optional[Any]						=  sd_pipe.prepare_inputs(snake_case_   )

    __a :  Optional[Any]						=  replicate(snake_case_   )
    __a :  Optional[int]						=  shard(snake_case_   )

    __a :  int						=  jax.random.PRNGKey(0   )
    __a :  str						=  jax.random.split(snake_case_					,				jax.device_count()   )

    __a :  int						=  sd_pipe(snake_case_					,				snake_case_					,				snake_case_					,				num_inference_steps=2_5					,				jit=snake_case_   )[0]
    assert images.shape == (jax.device_count(), 1, 7_6_8, 7_6_8, 3)

    __a :  Optional[int]						=  images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]   )
    __a :  List[Any]						=  images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

    __a :  Dict						=  jnp.asarray(jax.device_get(image_slice.flatten()   )   )
    __a :  str						=  jnp.array([0.4238, 0.4414, 0.4395, 0.4453, 0.4629, 0.4590, 0.4531, 0.4_5508, 0.4512]   )
    print(f"output_slice: {output_slice}"   )
    assert jnp.abs(output_slice - expected_slice   ).max() < 1E-2







   def       lowerCAmelCase			(self :					Tuple   ):
    __a :  Optional[Any]						=  '''stabilityai/stable-diffusion-2'''
    __a				,       __a :  Dict						=  FlaxDPMSolverMultistepScheduler.from_pretrained(snake_case_					,				subfolder='''scheduler'''   )
    __a				,       __a :  List[Any]						=  FlaxStableDiffusionPipeline.from_pretrained(
        snake_case_					,				scheduler=snake_case_					,				revision='''bf16'''					,				dtype=jnp.bfloataa					,				)
    __a :  Union[str, Any]						=  scheduler_params

    __a :  List[Any]						=  '''A painting of a squirrel eating a burger'''
    __a :  Any						=  jax.device_count()
    __a :  Any						=  num_samples * [prompt]
    __a :  List[Any]						=  sd_pipe.prepare_inputs(snake_case_   )

    __a :  Tuple						=  replicate(snake_case_   )
    __a :  Dict						=  shard(snake_case_   )

    __a :  Dict						=  jax.random.PRNGKey(0   )
    __a :  Dict						=  jax.random.split(snake_case_					,				jax.device_count()   )

    __a :  str						=  sd_pipe(snake_case_					,				snake_case_					,				snake_case_					,				num_inference_steps=2_5					,				jit=snake_case_   )[0]
    assert images.shape == (jax.device_count(), 1, 7_6_8, 7_6_8, 3)

    __a :  List[Any]						=  images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]   )
    __a :  Any						=  images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

    __a :  List[Any]						=  jnp.asarray(jax.device_get(image_slice.flatten()   )   )
    __a :  List[Any]						=  jnp.array([0.4336, 0.4_2969, 0.4453, 0.4199, 0.4297, 0.4531, 0.4434, 0.4434, 0.4297]   )
    print(f"output_slice: {output_slice}"   )
    assert jnp.abs(output_slice - expected_slice   ).max() < 1E-2
