File size: 1,544 Bytes
4121c7d |
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 |
import gradio as gr
from transformers import pipeline
from flask import Flask, request, jsonify
import os
from huggingface_hub import login
# Retrieve the Hugging Face token from environment variables
HF_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
if HF_TOKEN is None:
raise ValueError("Hugging Face token is not set in environment variables")
# Authenticate with Hugging Face
try:
login(token=HF_TOKEN, add_to_git_credential=True)
except ValueError as e:
print(f"Error during login: {e}")
raise
# Load your model and tokenizer
model_id = "rish13/llm_for_advanced_materials" # Replace with your model repo ID
model = pipeline('text-generation', model=model_id)
# Define Gradio interface
def generate_text(prompt):
return model(prompt)[0]['generated_text']
gradio_interface = gr.Interface(fn=generate_text, inputs="text", outputs="text")
# Initialize Flask app
app = Flask(__name__)
@app.route('/search', methods=['POST'])
def predict_endpoint():
data = request.json
prompt = data.get('prompt', '')
generated_text = generate_text(prompt)
return jsonify({"result": generated_text})
@app.route('/')
def home():
return "Welcome to the text generation API. Use the /search endpoint to generate text."
if __name__ == "__main__":
# Run the Gradio interface in a separate thread
from threading import Thread
gradio_thread = Thread(target=lambda: gradio_interface.launch(share=False, inbrowser=True))
gradio_thread.start()
# Run the Flask app
app.run(host='0.0.0.0', port=5000)
|