| from app.logger import get_logger | |
| logger = get_logger(__name__) | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| import torch | |
| model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" | |
| # Load tokenizer and model | |
| tokenizer =AutoTokenizer.from_pretrained(model_id) | |
| model = AutoModelForCausalLM.from_pretrained(model_id) | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| model.to(device) | |
| def generate_story(caption: str, max_length: int = 256) -> str: | |
| logger.info("Generating story...") | |
| try: | |
| # Turn caption into a story prompt | |
| prompt = ( | |
| "<|system|>\n" | |
| "You are a helpful assistant.</s>\n" | |
| "<|user|>\n" | |
| f"Write a complete, short story about {caption}. Make sure the story has a clear ending.\n</s>\n" | |
| "<|assistant|>\n" | |
| ) | |
| # Tokenize and run through model | |
| inputs = tokenizer(prompt, return_tensors="pt").to(device) | |
| outputs = model.generate( | |
| **inputs, | |
| max_new_tokens=1000, | |
| do_sample=True, | |
| temperature=0.8, | |
| top_p=0.9, | |
| top_k=50, | |
| eos_token_id=tokenizer.eos_token_id | |
| ) | |
| # Decode and clean output | |
| generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| generated_story = generated_text[len(prompt):] # Strip prompt part | |
| return generated_story.replace(prompt, "").strip() | |
| except Exception as e: | |
| logger.exception(f"Failed to generate story: {str(e)}") | |
| raise | |