🚀 AstroLLaMA

AstroLLaMA 是一个专注于天文学和天体物理学领域的文本生成模型。它基于 Llama - 2 架构，使用来自 arXiv 的天文学摘要数据进行训练，能够在相关领域生成高质量的文本内容。你可以在我们的 Hugging Face 空间中体验该模型！

🚀 快速开始

加载模型

from transformers import AutoModelForCausalLM
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(
    pretrained_model_name_or_path="universeTBD/astrollama"
)
model = AutoModelForCausalLM.from_pretrained(
    pretrained_model_name_or_path="universeTBD/astrollama",
    device_map="auto",
)

根据提示生成文本

import torch
from transformers import pipeline

generator = pipeline(
    task="text-generation",
    model=model,
    tokenizer=tokenizer,
    device_map="auto"
)

# 示例提示，取自 https://arxiv.org/abs/2308.12823
prompt = "In this letter, we report the discovery of the highest redshift, " \
    "heavily obscured, radio-loud QSO candidate selected using JWST NIRCam/MIRI, " \
    "mid-IR, sub-mm, and radio imaging in the COSMOS-Web field. "

# 为保证结果可复现
torch.manual_seed(42)

generated_text = generator(
    prompt,
    do_sample=True,
    max_length=512
)

使用 AstroLLaMA 进行文本嵌入

texts = [
    "Abstract 1",
    "Abstract 2"
]
inputs = tokenizer(
    texts,
    return_tensors="pt",
    return_token_type_ids=False,
    padding=True,
    truncation=True,
    max_length=4096
)
inputs.to(model.device)
outputs = model(**inputs, output_hidden_states=True)

# 取隐藏状态的最后一层，计算所有标记的平均嵌入
embeddings = outputs["hidden_states"][-1][:, 1:, ...].mean(1).detach().cpu().numpy()

📄 许可证

本项目采用 MIT 许可证。

📦 模型信息

立即体验模型！ 点击此处 前往我们的 Hugging Face 空间。