Update main.py

main.py

@@ -3,24 +3,14 @@ from fastapi.responses import FileResponse
 from pydantic import BaseModel
 from fastapi import FastAPI
 
-import os
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-from transformers import GPT2LMHeadModel, GPT2Tokenizer, AutoModelForCausalLM, AutoTokenizer
-import torch
+model_name = "facebook/blenderbot-400M-distill"
+# facebook/blenderbot-400M-distill
 
 app = FastAPI()
-name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-customGen = False
-
-# microsoft/DialoGPT-small
-# microsoft/DialoGPT-medium
-# microsoft/DialoGPT-large
-
-# mistralai/Mixtral-8x7B-Instruct-v0.1
-
-# Load the Hugging Face GPT-2 model and tokenizer
-model = AutoModelForCausalLM.from_pretrained(name)
-tokenizer = AutoTokenizer.from_pretrained(name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 class req(BaseModel):
   prompt: str
@@ -34,6 +24,7 @@ def read_root():
 def read_root(data: req):
   print("Prompt:", data.prompt)
   print("Length:", data.length)
+  print("Generating")
 
   input_text = data.prompt