GeoLLM / Task2 /LLM.py
Ciallo0d00's picture
Upload folder using huggingface_hub
badcf3c verified
from openai import OpenAI
import os
import pandas as pd
import google.generativeai as genai
import requests
# 输入模型系列和内容
def zero_shot(model_series, model_name, content):
# 读取txt文件
df = pd.read_csv('./LLM_APIs.txt')
# 由于列名包含制表符,需要先分割列名
df = pd.DataFrame([x.split('\t') for x in df.values.flatten()], columns=['name','series','API'])
# 获取对应的API密钥
api_key = df[df['name'] == model_series]['API'].values[0]
# 根据不同模型系列调用不同API
if model_series == 'gpt':
client = OpenAI(
api_key=api_key,
base_url="https://api.bianxie.ai/v1"
)
# print("##1")
completion = client.chat.completions.create(
model = model_name,
messages=[
{
"role": "user",
"content": content
}
]
)
# print("##2")
return completion.choices[0].message
elif model_series == 'llama':
client = OpenAI(
api_key=api_key,
base_url="https://api.bianxie.ai/v1"
)
completion = client.chat.completions.create(
model=model_name,
messages=[
{
"role": "user",
"content": content
}
]
)
return completion.choices[0].message
elif model_series == 'deepSeek':
client = OpenAI(api_key=api_key,
base_url="https://api.studio.nebius.ai/v1")
response = client.chat.completions.create(
model=model_name,
messages=[
{"role": "user", "content": content},
],
stream=False
)
# client = OpenAI(api_key=api_key,
# base_url="https://api.deepseek.com")
# response = client.chat.completions.create(
# model=model_name,
# messages=[
# {"role": "user", "content": content},
# ],
# stream=False
# )
return response.choices[0].message
elif model_series == 'gemini':
# genai.configure(api_key=api_key)
# model = genai.GenerativeModel(model_name)
# response = model.generate_content(content)
# print(response.text)
# return response.text
api_key = api_key
url = 'https://api.bianxie.ai/v1/chat/completions'
headers = {
'Content-Type': 'application/json',
'Authorization': f'Bearer {api_key}'
}
data = {
'model': model_name,
'messages': [{'role': 'user', 'content': content}],
}
response = requests.post(url, headers=headers, json=data)
content = response.json()['choices'][0]['message']
return content
elif model_series == 'claude':
api_key = api_key
url = 'https://api.bianxie.ai/v1/chat/completions'
headers = {
'Content-Type': 'application/json',
'Authorization': f'Bearer {api_key}'
}
data = {
'model': model_name,
'messages': [{'role': 'user', 'content': content}],
}
response = requests.post(url, headers=headers, json=data)
print(response.json())
content = response.json()['choices'][0]['message']
return content
else:
return "不支持的模型系列"