|
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 |
|
|