File size: 2,448 Bytes
08aad81 a7c32b2 2fd3a49 16ce850 a7c32b2 74c9bed 2fd3a49 a7c32b2 08aad81 2fd3a49 a7c32b2 d28821f 642d8b4 bf28dd1 2754c6f bf28dd1 08aad81 a7c32b2 af440aa a7c32b2 d28821f a7c32b2 d28821f a7c32b2 bf28dd1 2fd3a49 |
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 |
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import gradio as gr
import requests
import threading
app = FastAPI()
# Load model and tokenizer once
tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
# In-memory chat history by user
chat_history = {}
@app.get("/")
async def root():
return {"message": "🟢 API is running. Use /ai?query=Hello&user_id=yourname"}
@app.get("/ai")
async def chat(request: Request):
query_params = dict(request.query_params)
user_input = query_params.get("query", "")
user_id = query_params.get("user_id", "default")
if not user_input:
return JSONResponse({"error": "Missing 'query' parameter"}, status_code=400)
new_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt')
user_history = chat_history.get(user_id, [])
bot_input_ids = torch.cat(user_history + [new_input_ids], dim=-1) if user_history else new_input_ids
output_ids = model.generate(bot_input_ids, max_new_tokens=100, pad_token_id=tokenizer.eos_token_id)
response = tokenizer.decode(output_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
chat_history[user_id] = [bot_input_ids, output_ids]
return JSONResponse({"reply": response})
# Gradio UI to call your /ai endpoint easily via browser
def gradio_chat(user_input, user_id="default"):
if not user_input:
return "Please enter some text."
url = f"https://Trigger82--API.hf.space/ai?query={user_input}&user_id={user_id}"
try:
res = requests.get(url)
if res.status_code == 200:
return res.json().get("reply", "No reply")
return f"Error: {res.status_code}"
except Exception as e:
return f"Exception: {e}"
iface = gr.Interface(
fn=gradio_chat,
inputs=[gr.Textbox(label="Your Message"), gr.Textbox(label="User ID", value="default")],
outputs="text",
title="Chat with DialoGPT API",
description="Type your message and user id to chat with the model."
)
# Launch Gradio app in a thread alongside FastAPI
def run_gradio():
iface.launch(server_name="0.0.0.0", server_port=7861, share=False)
threading.Thread(target=run_gradio).start()
# No need for uvicorn.run here on Spaces; it manages startup automatically |