feat: Create OllamaGenerator class for Pelican markdown generation
This commit is contained in:
parent
858f7c8d0a
commit
675c7d1701
28
ollama_generator.py
Normal file
28
ollama_generator.py
Normal file
@ -0,0 +1,28 @@
|
||||
import os
|
||||
import requests
|
||||
|
||||
class OllamaGenerator:
|
||||
def __init__(self, title: str, content: str):
|
||||
self.title = title
|
||||
self.content = content
|
||||
self.ollama_url = f"{os.getenv('OLLAMA_PROTOCOL')}://{os.getenv('OLLAMA_HOST')}:{os.getenv('OLLAMA_PORT')}/api/generate"
|
||||
|
||||
def generate_markdown(self) -> str:
|
||||
data = {
|
||||
"prompt": f"Title: {self.title}\n\nContent:\n{self.content}",
|
||||
"max_tokens": 1000,
|
||||
"temperature": 0.7,
|
||||
"top_p": 0.9,
|
||||
"presence_penalty": 0,
|
||||
"frequency_penalty": 0
|
||||
}
|
||||
|
||||
response = requests.post(self.ollama_url, json=data)
|
||||
if response.status_code == 200:
|
||||
return response.json()["choices"][0]["text"]
|
||||
else:
|
||||
raise Exception(f"Failed to generate markdown: {response.text}")
|
||||
|
||||
def save_to_file(self, filename: str) -> None:
|
||||
with open(filename, "w") as f:
|
||||
f.write(self.generate_markdown())
|
Loading…
x
Reference in New Issue
Block a user