HusseinBashir's picture
Update app.py
2546b1c verified
raw
history blame contribute delete
600 Bytes
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from fastapi import FastAPI
from pydantic import BaseModel
import torch
tokenizer = AutoTokenizer.from_pretrained("cmlang/somali-flan-t5")
model = AutoModelForSeq2SeqLM.from_pretrained("cmlang/somali-flan-t5")
app = FastAPI()
class Message(BaseModel):
text: str
@app.post("/chatbot")
async def chatbot(msg: Message):
inputs = tokenizer(msg.text, return_tensors="pt")
outputs = model.generate(**inputs, max_length=150)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
return {"response": response}