Spaces:
Sleeping
Sleeping
File size: 2,523 Bytes
e84e505 | 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 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | """Provider wrappers for question generation.
Both providers expose complete(prompt) so dataset code can share one interface.
"""
import os
class DeepSeekSource:
"""DeepSeek-based question source."""
def __init__(self, model_name, temperature):
"""Initialize DeepSeek client from env key."""
from openai import OpenAI
api_key = os.environ.get("DEEPSEEK_API_KEY")
if not api_key:
raise SystemExit(
"DEEPSEEK_API_KEY is not set. Add it to .env. Get a key with "
"free credits at https://platform.deepseek.com (no card needed)."
)
self.client = OpenAI(api_key=api_key, base_url="https://api.deepseek.com")
self.model_name = model_name
self.temperature = temperature
self.label = "deepseek:" + model_name
def complete(self, prompt):
"""Send one prompt and return raw text."""
response = self.client.chat.completions.create(
model=self.model_name,
messages=[{"role": "user", "content": prompt}],
temperature=self.temperature,
response_format={"type": "json_object"},
)
return response.choices[0].message.content
class GeminiSource:
"""Gemini-based question source."""
def __init__(self, model_name, temperature):
"""Initialize Gemini client from env key."""
from google import genai
api_key = os.environ.get("GEMINI_API_KEY")
if not api_key:
raise SystemExit("GEMINI_API_KEY is not set. Add it to .env.")
self.client = genai.Client(api_key=api_key)
self.model_name = model_name
self.temperature = temperature
self.label = "gemini:" + model_name
def complete(self, prompt):
"""Send one prompt and return raw text."""
from google.genai import types
response = self.client.models.generate_content(
model=self.model_name,
contents=prompt,
config=types.GenerateContentConfig(
response_mime_type="application/json",
temperature=self.temperature,
),
)
return response.text
def build_source(provider, model_name, temperature):
"""Return provider-specific source instance."""
if provider == "deepseek":
return DeepSeekSource(model_name, temperature)
if provider == "gemini":
return GeminiSource(model_name, temperature)
raise ValueError(f"Unknown provider: {provider}")
|