app.py

from transformers import AutoModelForCausalLM, AutoTokenizer

# Load the Mistral model and tokenizer from Hugging Face Hub
model_name = "mistralai/mistral-7b"  # Change this to the desired model name
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Function to generate text
def generate_text(prompt: str, max_length: int = 50):
    # Tokenize input prompt
    inputs = tokenizer(prompt, return_tensors="pt")
    
    # Generate text
    outputs = model.generate(
        input_ids=inputs['input_ids'],
        max_length=max_length,  # Max tokens in output
        num_return_sequences=1,  # Number of sequences to generate
        temperature=0.7,  # Adjust for randomness (higher = more random)
        top_p=0.9,  # Use top-p sampling for diversity
        top_k=50  # Top-k sampling for diversity
    )
    
    # Decode the generated text and return it
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return generated_text

# Example usage
prompt = "The future of artificial intelligence is"
generated_output = generate_text(prompt)
print("Generated Text:", generated_output)