Spaces:
Running
Running
File size: 789 Bytes
809141b 6bcf7d9 809141b b8ac549 809141b b8ac549 809141b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 |
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer, pipeline
import torch
class HuggingFaceAgent:
def __init__(self, model_id, access_token):
self.tokenizer = AutoTokenizer.from_pretrained(model_id, token=access_token)
self.model = AutoModelForCausalLM.from_pretrained(
model_id,
token=access_token,
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
device_map="auto"
)
self.pipeline = pipeline("text-generation", model=self.model, tokenizer=self.tokenizer)
def responder(self, prompt):
respuesta = self.pipeline(prompt, max_new_tokens=256, do_sample=True, temperature=0.7)[0]["generated_text"]
return respuesta[len(prompt):].strip()
|