models/tiny_aya.py

from huggingface_hub import InferenceClient

class TinyAya:
    def __init__(self, token):
        self.client = InferenceClient(token=token)
        self.model_id = "CohereLabs/tiny-aya-global"

    def generate_stream(self, prompt, max_tokens=1500, temperature=0.1):
        try:
            for message in self.client.chat_completion(
                model=self.model_id,
                messages=[{"role": "user", "content": prompt}],
                max_tokens=max_tokens,
                temperature=temperature,
                stream=True,
            ):
                if message.choices:
                    content = message.choices[0].delta.content
                    if content:
                        yield content
        except Exception as e:
            yield f" TinyAya Error: {e}"

    def generate(self, prompt, max_tokens=500, temperature=0.1):
        return "".join(self.generate_stream(prompt, max_tokens=max_tokens, temperature=temperature))