File size: 1,221 Bytes
6f15a2e
c654f8e
6f15a2e
 
1232177
 
 
 
73e820c
 
6f15a2e
 
73e820c
 
 
 
059c9d2
c654f8e
73e820c
6f15a2e
 
1232177
6f15a2e
059c9d2
784fe5f
1232177
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
import os
from transformers import AutoTokenizer, AutoModelForCausalLM  # Ensure correct model class

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  # Set pad_token to eos_token
        model = AutoModelForCausalLM.from_pretrained(model_name)  # Ensure correct model class
        print("Model and tokenizer loaded successfully.")
    return tokenizer, model

async def process_text_local(model_name, text):
    tokenizer, model = load_model(model_name)
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)  # Set max_length to 512
    outputs = model.generate(**inputs, max_length=512)
    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return result