| import os |
| from transformers import AutoTokenizer, AutoModelForCausalLM |
|
|
| HUGGINGFACE_API_KEY = os.getenv("HUGGINGFACE_API_KEY") |
| if HUGGINGFACE_API_KEY is None: |
| raise ValueError("Hugging Face API key is not set. Please add it as a secret in your Hugging Face Space settings.") |
| print(f"Using Hugging Face API Key: {HUGGINGFACE_API_KEY}") |
|
|
| model = None |
| tokenizer = None |
|
|
| def load_model(model_name): |
| global tokenizer, model |
| if not tokenizer or not model: |
| print("Loading model and tokenizer...") |
| tokenizer = AutoTokenizer.from_pretrained(model_name) |
| tokenizer.pad_token = tokenizer.eos_token |
| model = AutoModelForCausalLM.from_pretrained(model_name) |
| print("Model and tokenizer loaded successfully.") |
| return tokenizer, model |
|
|
| async def process_text_local(model_name, text): |
| print("Loading model and tokenizer...") |
| tokenizer, model = load_model(model_name) |
| print("Encoding text...") |
| inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512) |
| print("Text encoded successfully.") |
| print("Generating output...") |
| outputs = model.generate(**inputs, max_length=512) |
| print("Output generated successfully.") |
| result = tokenizer.decode(outputs[0], skip_special_tokens=True) |
| print("Output decoded successfully.") |
| return result |
|
|