File size: 1,244 Bytes
4d1930d
702c209
 
 
d07fe3b
702c209
d07fe3b
21eb127
937e294
 
b248a5f
937e294
 
702c209
937e294
 
702c209
937e294
702c209
 
d07fe3b
4d1930d
702c209
4d1930d
 
 
937e294
4d1930d
 
 
 
 
 
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
import argparse
import os
import requests
from dotenv import load_dotenv

load_dotenv()  # Load environment variables from .env

def generate_text(prompt):
    # Use the correct model URL for EleutherAI's GPT-J 6B
    url = "https://api-inference.huggingface.co/models/EleutherAI/gpt-j-6b"
    headers = {"Authorization": f"Bearer {os.getenv('HUGGINGFACE_API_KEY')}"}
    
    # Send the request with the prompt as input
    response = requests.post(url, headers=headers, json={"inputs": prompt})
    
    # Check the response status and return the generated text if successful
    if response.status_code == 200:
        return response.json()[0]["generated_text"]  # Adjust for API's JSON structure
    else:
        return f"Error: {response.status_code} - {response.text}"

def cli_interface():
    parser = argparse.ArgumentParser(description="Command-line interaction with the deployed model.")
    parser.add_argument("--task", type=str, help="The prompt or command to generate text for")
    args = parser.parse_args()

    # Use the provided prompt or default to "Tell me a joke"
    task = args.task if args.task else "Tell me a joke"
    result = generate_text(task)
    print(result)

if __name__ == "__main__":
    cli_interface()